Fix(be): Avoid decompression race condition (#718)

Author: robTheBuildr

src/litdata/streaming/config.py

@@ -14,9 +14,10 @@
 import logging
 import os
 from collections import defaultdict
+from time import sleep, time
 from typing import Any, Optional
 
-from litdata.constants import _INDEX_FILENAME
+from litdata.constants import _INDEX_FILENAME, _MAX_WAIT_TIME
 from litdata.debugger import ChromeTraceColors, _get_log_msg
 from litdata.streaming.compression import _COMPRESSORS, Compressor
 from litdata.streaming.downloader import get_downloader
@@ -183,6 +184,21 @@ def try_decompress(self, local_chunkpath: str) -> None:
         if os.path.exists(target_local_chunkpath):
             return
 
+        # Ensure that the compressed file exists and is fully downloaded
+        start_time = time()
+        assert self._chunks is not None
+
+        filename = os.path.basename(local_chunkpath)
+        chunk_index = self._get_chunk_index_from_filename(filename)
+        chunk_bytes = self._chunks[chunk_index]["chunk_size"]
+        exists = os.path.exists(local_chunkpath) and os.stat(local_chunkpath).st_size >= chunk_bytes
+        while not exists:
+            sleep(0.1)
+            exists = os.path.exists(local_chunkpath) and os.stat(local_chunkpath).st_size >= chunk_bytes
+
+            if (time() - start_time) > _MAX_WAIT_TIME:
+                raise FileNotFoundError(f"The {local_chunkpath} hasn't been found.")
+
         with open(local_chunkpath, "rb") as f:
             data = f.read()