Fix errors when using compression and r2 in optimize() (#715)

Author: pwgardipee

src/litdata/processing/data_processor.py

@@ -217,6 +217,7 @@ def keep_path(path: str) -> bool:
         "s3_connections",
         "s3_folders",
         "snowflake_connections",
+        "lightning_storage",
     ]
     return all(p not in path for p in paths)
 

src/litdata/streaming/dataset.py

@@ -194,6 +194,12 @@ def __init__(
         self.num_workers: int = 1
         self.batch_size: int = 1
         self._encryption = encryption
+        # Ensure data_connection_id is included in storage_options if available from input_dir
+        if input_dir.data_connection_id and storage_options is not None:
+            storage_options = storage_options.copy()
+            storage_options["data_connection_id"] = input_dir.data_connection_id
+        elif input_dir.data_connection_id and storage_options is None:
+            storage_options = {"data_connection_id": input_dir.data_connection_id}
         self.storage_options = storage_options
         self.session_options = session_options
         self.max_pre_download = max_pre_download

src/litdata/streaming/resolver.py

@@ -62,7 +62,7 @@ def _resolve_dir(dir_path: Optional[Union[str, Path, Dir]]) -> Dir:
         raise ValueError(f"`dir_path` must be either a string, Path, or Dir, got: {type(dir_path)}")
 
     if isinstance(dir_path, str):
-        cloud_prefixes = ("s3://", "gs://", "azure://", "hf://")
+        cloud_prefixes = ("s3://", "gs://", "r2://", "azure://", "hf://")
         if dir_path.startswith(cloud_prefixes):
             return Dir(path=None, url=dir_path)
 

src/litdata/utilities/dataset_utilities.py

@@ -84,7 +84,12 @@ def subsample_streaming_dataset(
         if index_path is not None:
             copy_index_to_cache_index_filepath(index_path, cache_index_filepath)
         else:
-            downloader = get_downloader(input_dir.url, input_dir.path, [], storage_options, session_options)
+            # 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:
@@ -159,6 +164,7 @@ def _should_replace_path(path: Optional[str]) -> bool:
         or path.startswith("/teamspace/s3_folders/")
         or path.startswith("/teamspace/gcs_folders/")
         or path.startswith("/teamspace/gcs_connections/")
+        or path.startswith("/teamspace/lightning_storage/")
     )
 
 

tests/streaming/test_dataset.py

@@ -867,6 +867,24 @@ def fn(remote_chunkpath: str, local_chunkpath: str):
     )  # it won't be None, and a cache dir will be created
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="Not tested on windows and MacOs")
+def test_r2_streaming_dataset(monkeypatch, tmpdir):
+    """Test that data_connection_id is properly merged into storage_options."""
+    downloader = mock.MagicMock()
+
+    def fn(remote_chunkpath: str, local_chunkpath: str):
+        with open(local_chunkpath, "w") as f:
+            json.dump({"chunks": [{"chunk_size": 2, "filename": "0.bin"}]}, f)
+
+    downloader.download_file = fn
+
+    monkeypatch.setattr(dataset_utilities_module, "get_downloader", mock.MagicMock(return_value=downloader))
+
+    dataset = StreamingDataset(input_dir="r2://random_bucket/optimized_tiny_imagenet")
+    assert dataset.input_dir.url == "r2://random_bucket/optimized_tiny_imagenet"
+    assert dataset.input_dir.path.endswith("chunks/9537e9e392ad87a4d38d05dfe28c329a/9537e9e392ad87a4d38d05dfe28c329a")
+
+
 class EmulateS3StreamingDataset(StreamingDataset):
     def _create_cache(self, worker_env: _WorkerEnv) -> Cache:
         cache_dir = os.path.join(self.input_dir.path)

tests/utilities/test_dataset_utilities.py

@@ -27,6 +27,7 @@ def test_should_replace_path():
     assert _should_replace_path("/teamspace/s3_folders/...")
     assert _should_replace_path("/teamspace/gcs_folders/...")
     assert _should_replace_path("/teamspace/gcs_connections/...")
+    assert _should_replace_path("/teamspace/lightning_storage/...")
     assert not _should_replace_path("something_else")