From ed371c090794167bfcdadb200c17dbe1e17efacb Mon Sep 17 00:00:00 2001 From: Javier de Jesus Date: Mon, 22 Jun 2026 01:26:32 +0000 Subject: [PATCH] fix: honor index_path for local input_dir in subsample_streaming_dataset The index_path copy was gated behind isinstance(input_dir.url, str), so a user-provided index_path was silently ignored when input_dir is a local directory (url is None), raising a misleading "doesn't contain any index.json file" error. Copy the provided index into the local cache regardless of whether input_dir.url is a string. --- src/litdata/utilities/dataset_utilities.py | 22 +++++++++++----------- tests/utilities/test_dataset_utilities.py | 21 +++++++++++++++++++++ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/src/litdata/utilities/dataset_utilities.py b/src/litdata/utilities/dataset_utilities.py index 6aa3d154f..5d12c460e 100644 --- a/src/litdata/utilities/dataset_utilities.py +++ b/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) diff --git a/tests/utilities/test_dataset_utilities.py b/tests/utilities/test_dataset_utilities.py index 63a9f4688..d53e40bd0 100644 --- a/tests/utilities/test_dataset_utilities.py +++ b/tests/utilities/test_dataset_utilities.py @@ -7,6 +7,7 @@ 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, @@ -14,6 +15,7 @@ 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