Remove unnecessary fixed sleep by adding predicate-based path check (Lightning-AI#700)

* Remove unnecessary fixed sleep by adding predicate-based path check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update dataset_utilities.py

fixed checks

* wip

* Update src/litdata/utilities/dataset_utilities.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* fix:last commit

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: bhimrazy <bhimrajyadav977@gmail.com>

Author: 4 people

src/litdata/utilities/dataset_utilities.py

@@ -5,7 +5,7 @@
 import shutil
 import tempfile
 import time
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import numpy as np
 
@@ -16,6 +16,27 @@
 from litdata.utilities.subsample import shuffle_lists_together, subsample_filenames_and_roi
 
 
+def wait_for_predicate(
+    predicate: Callable[[], bool],
+    timeout: float,
+) -> bool:
+    """Wait until the given predicate becomes True or the timeout expires.
+
+    Args:
+        predicate: A function returning a boolean condition to check.
+        timeout: Maximum time (in seconds) to wait.
+
+    Returns:
+        True if predicate became True within timeout, else False.
+    """
+    start = time.time()
+    while time.time() - start < timeout:
+        if predicate():
+            return True
+        time.sleep(0.01)
+    return False
+
+
 def subsample_streaming_dataset(
     input_dir: Dir,
     cache_dir: Optional[Dir] = None,
@@ -66,9 +87,10 @@ def subsample_streaming_dataset(
             downloader = get_downloader(input_dir.url, input_dir.path, [], storage_options, session_options)
             downloader.download_file(os.path.join(input_dir.url, _INDEX_FILENAME), cache_index_filepath)
 
-    time.sleep(0.5)  # Give some time for the file to be available
+    def path_exists(p: str) -> bool:
+        return wait_for_predicate(lambda: os.path.exists(p), timeout=0.5)
 
-    if not os.path.exists(input_dir.path):
+    if not path_exists(input_dir.path):
         raise FileNotFoundError(f"The provided dataset path `{input_dir.path}` does not exist.")
 
     if os.path.exists(os.path.join(input_dir.path, _INDEX_FILENAME)):