Resolve path for StreamingDataset (#19094)

Co-authored-by: thomas <[email protected]>

Author: tchaton

src/lightning/data/streaming/dataset.py

@@ -110,7 +110,7 @@ def __init__(
     def _create_cache(self, worker_env: _WorkerEnv) -> Cache:
         env = Environment(dist_env=self.distributed_env, worker_env=worker_env)
 
-        if self.input_dir.path is None:
+        if _should_replace_path(self.input_dir.path):
             cache_path = _try_create_cache_dir(
                 input_dir=self.input_dir.path if self.input_dir.path else self.input_dir.url, shard_rank=env.shard_rank
             )
@@ -427,6 +427,13 @@ def _collect_distributed_state_dict(state_dict: Dict[str, Any], world_size: int)
     return state_dict_out
 
 
+def _should_replace_path(path: str) -> bool:
+    if path is None or path == "":
+        return True
+
+    return "/datasets/" in path or "_connections/" in path
+
+
 def _is_in_dataloader_worker() -> bool:
     return get_worker_info() is not None
 

tests/tests_data/streaming/test_dataset.py

@@ -24,7 +24,7 @@
 from lightning import seed_everything
 from lightning.data.streaming import Cache, functions
 from lightning.data.streaming.constants import _TIME_FORMAT
-from lightning.data.streaming.dataset import StreamingDataset, _try_create_cache_dir
+from lightning.data.streaming.dataset import StreamingDataset, _should_replace_path, _try_create_cache_dir
 from lightning.data.streaming.item_loader import TokensLoader
 from lightning.data.streaming.shuffle import FullShuffle, NoShuffle
 from lightning.data.utilities.env import _DistributedEnv
@@ -59,6 +59,14 @@ def test_streaming_dataset(tmpdir, monkeypatch):
     assert len(dataloader) == 6
 
 
+def test_should_replace_path():
+    assert _should_replace_path(None)
+    assert _should_replace_path("")
+    assert _should_replace_path(".../datasets/...")
+    assert _should_replace_path(".../_connections/...")
+    assert not _should_replace_path("something_else")
+
+
 @pytest.mark.parametrize("drop_last", [False, True])
 def test_streaming_dataset_distributed_no_shuffle(drop_last, tmpdir):
     seed_everything(42)