Streaming Dataset: Resolve chunks eviction (#19214)

Co-authored-by: thomas <[email protected]>

Author: tchaton

src/lightning/data/streaming/reader.py

@@ -68,6 +68,7 @@ def __init__(
 
         # FIXME: This should be divided by the number of nodes to provide a more granular support with scaling out
         self._delete_chunks_when_processed = self._config.num_bytes > max_cache_size if max_cache_size else False
+        self._has_exited = False
 
     def download(self, chunk_indexes: List[int]) -> None:
         """Receive the list of the chunk indices to download for the current epoch."""
@@ -111,7 +112,7 @@ def _maybe_delete_chunks(self) -> None:
 
     def _can_delete_chunk(self) -> bool:
         if self._delete_chunks_when_processed:
-            return self._pre_download_counter == self._max_pre_download - 1
+            return self._pre_download_counter >= self._max_pre_download - 1
         return self._max_cache_size is not None and _get_folder_size(self._parent_cache_dir) >= self._max_cache_size
 
     def _pre_load_chunk(self, chunk_index: int) -> None:
@@ -120,9 +121,10 @@ def _pre_load_chunk(self, chunk_index: int) -> None:
 
     def run(self) -> None:
         while True:
-            if self._pre_download_counter <= self._max_pre_download:
+            if self._pre_download_counter < self._max_pre_download:
                 chunk_index = _get_from_queue(self._to_download_queue)
                 if chunk_index == _END_TOKEN:
+                    self._has_exited = True
                     return
 
                 if chunk_index is not None:

tests/tests_data/streaming/test_reader.py

@@ -1,11 +1,13 @@
 import os
 import shutil
+from time import sleep
 
 import numpy as np
+from lightning.data.streaming import reader
 from lightning.data.streaming.cache import Cache
 from lightning.data.streaming.config import ChunkedIndex
 from lightning.data.streaming.item_loader import PyTreeLoader
-from lightning.data.streaming.reader import PrepareChunksThread, _get_folder_size
+from lightning.data.streaming.reader import _END_TOKEN, PrepareChunksThread, _get_folder_size
 from lightning_cloud.resolver import Dir
 
 
@@ -36,40 +38,11 @@ def test_reader_chunk_removal(tmpdir):
     shutil.rmtree(cache_dir)
     os.makedirs(cache_dir, exist_ok=True)
 
-    generated = []
     for i in range(25):
-        generated.append([i, len(os.listdir(cache_dir))])
+        assert len(os.listdir(cache_dir)) <= 3
         index = ChunkedIndex(i, cache._get_chunk_index_from_index(i), is_last_index=i == 24)
         assert cache[index] == i
 
-    assert generated == [
-        [0, 0],
-        [1, 2],
-        [2, 2],
-        [3, 3],
-        [4, 3],
-        [5, 3],
-        [6, 3],
-        [7, 3],
-        [8, 3],
-        [9, 3],
-        [10, 3],
-        [11, 3],
-        [12, 3],
-        [13, 3],
-        [14, 3],
-        [15, 3],
-        [16, 3],
-        [17, 3],
-        [18, 3],
-        [19, 3],
-        [20, 3],
-        [21, 3],
-        [22, 3],
-        [23, 3],
-        [24, 3],
-    ]
-
     assert len(os.listdir(cache_dir)) == 3
 
 
@@ -82,7 +55,9 @@ def test_get_folder_size(tmpdir):
     assert _get_folder_size(tmpdir) == 928 * 2
 
 
-def test_prepare_chunks_thread(tmpdir):
+def test_prepare_chunks_thread_eviction(tmpdir, monkeypatch):
+    monkeypatch.setattr(reader, "_LONG_DEFAULT_TIMEOUT", 0.1)
+
     cache_dir = os.path.join(tmpdir, "cache_dir")
     os.makedirs(cache_dir, exist_ok=True)
     cache = Cache(input_dir=cache_dir, chunk_size=2, max_cache_size=28020)
@@ -95,8 +70,32 @@ def test_prepare_chunks_thread(tmpdir):
 
     cache._reader._try_load_config()
 
-    thread = PrepareChunksThread(cache._reader.config, item_loader=PyTreeLoader(), max_cache_size=1)
-    assert thread._delete_chunks_when_processed
+    assert len(os.listdir(cache_dir)) == 14
 
     thread = PrepareChunksThread(cache._reader.config, item_loader=PyTreeLoader(), max_cache_size=10000)
     assert not thread._delete_chunks_when_processed
+
+    thread = PrepareChunksThread(cache._reader.config, item_loader=PyTreeLoader(), max_cache_size=1)
+    assert thread._delete_chunks_when_processed
+
+    thread.start()
+
+    assert thread._pre_download_counter == 0
+
+    thread.download([0, 1, 2, 3, 4, 5, _END_TOKEN])
+
+    while thread._pre_download_counter == 0:
+        sleep(0.01)
+
+    assert not thread._has_exited
+
+    for i in range(5):
+        thread.delete([i])
+        while len(os.listdir(cache_dir)) != 14 - (i + 1):
+            sleep(0.01)
+
+    assert thread._pre_download_counter <= 2
+
+    assert len(os.listdir(cache_dir)) == 9
+    assert thread._has_exited
+    thread.join()