fix: honor index_path when input_dir is a local directory

src/litdata/utilities/dataset_utilities.py

@@ -80,18 +80,18 @@ def subsample_streaming_dataset(
     cache_index_filepath = os.path.join(input_dir.path, _INDEX_FILENAME)
 
     # Check if `index.json` file exists in cache path
-    if not os.path.exists(cache_index_filepath) and isinstance(input_dir.url, str):
+    if not os.path.exists(cache_index_filepath) and index_path is not None:
+        # A user-provided `index_path` should be honored for both local and remote `input_dir`.
+        copy_index_to_cache_index_filepath(index_path, cache_index_filepath)
+    elif not os.path.exists(cache_index_filepath) and isinstance(input_dir.url, str):
         assert input_dir.url is not None
-        if index_path is not None:
-            copy_index_to_cache_index_filepath(index_path, cache_index_filepath)
-        else:
-            # Merge data_connection_id from resolved directory into storage_options for R2 connections
-            merged_storage_options = storage_options.copy() if storage_options is not None else {}
-            if hasattr(input_dir, "data_connection_id") and input_dir.data_connection_id:
-                merged_storage_options["data_connection_id"] = input_dir.data_connection_id
-
-            downloader = get_downloader(input_dir.url, input_dir.path, [], merged_storage_options, session_options)
-            downloader.download_file(os.path.join(input_dir.url, _INDEX_FILENAME), cache_index_filepath)
+        # Merge data_connection_id from resolved directory into storage_options for R2 connections
+        merged_storage_options = storage_options.copy() if storage_options is not None else {}
+        if hasattr(input_dir, "data_connection_id") and input_dir.data_connection_id:
+            merged_storage_options["data_connection_id"] = input_dir.data_connection_id
+
+        downloader = get_downloader(input_dir.url, input_dir.path, [], merged_storage_options, session_options)
+        downloader.download_file(os.path.join(input_dir.url, _INDEX_FILENAME), cache_index_filepath)
 
     def path_exists(p: str) -> bool:
         return wait_for_predicate(lambda: os.path.exists(p), timeout=0.5)

tests/utilities/test_dataset_utilities.py

@@ -7,13 +7,15 @@
 import litdata.utilities
 import litdata.utilities.dataset_utilities
 from litdata.constants import _DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR, _INDEX_FILENAME
+from litdata.streaming.resolver import Dir
 from litdata.utilities.dataset_utilities import (
     _should_replace_path,
     _try_create_cache_dir,
     adapt_mds_shards_to_chunks,
     generate_roi,
     get_default_cache_dir,
     load_index_file,
+    subsample_streaming_dataset,
 )
 
 
@@ -99,6 +101,25 @@ def test_adapt_mds_shards_to_chunks(mosaic_mds_index_data):
     assert len(mosaic_mds_index_data["shards"]) == len(adapted_data["chunks"])
 
 
+def test_subsample_streaming_dataset_with_local_index_path(tmpdir):
+    # `input_dir` is a local directory that does NOT contain an index.json file.
+    data_dir = tmpdir.mkdir("data")
+    # The index.json lives in a separate directory provided via `index_path`.
+    index_dir = tmpdir.mkdir("index")
+    index_data = {"chunks": [{"chunk_size": 30, "filename": "chunk-0.bin"}], "config": {}}
+    with open(os.path.join(index_dir, _INDEX_FILENAME), "w") as f:
+        f.write(json.dumps(index_data))
+
+    input_dir = Dir(path=str(data_dir), url=None)
+    with mock.patch.dict(os.environ, {}, clear=True):
+        subsampled_files, roi = subsample_streaming_dataset(input_dir, index_path=str(index_dir))
+
+    assert subsampled_files == ["chunk-0.bin"]
+    assert roi == [(0, 30)]
+    # The user-provided index is materialized into the local dataset directory.
+    assert os.path.exists(os.path.join(str(data_dir), _INDEX_FILENAME))
+
+
 def test_get_default_cache_dir():
     with mock.patch.dict(os.environ, {}, clear=True):
         assert get_default_cache_dir() == _DEFAULT_CACHE_DIR