Add support for resolving directories in /teamspace/lightning_storage

src/litdata/streaming/resolver.py

@@ -89,6 +89,9 @@ def _resolve_dir(dir_path: Optional[Union[str, Path, Dir]]) -> Dir:
     if dir_path_absolute.startswith("/teamspace/gcs_folders") and len(dir_path_absolute.split("/")) > 3:
         return _resolve_gcs_folders(dir_path_absolute)
 
+    if dir_path_absolute.startswith("/teamspace/lightning_storage") and len(dir_path_absolute.split("/")) > 3:
+        return _resolve_lightning_storage(dir_path_absolute)
+
     if dir_path_absolute.startswith("/teamspace/datasets") and len(dir_path_absolute.split("/")) > 3:
         return _resolve_datasets(dir_path_absolute)
 
@@ -246,6 +249,28 @@ def _resolve_gcs_folders(dir_path: str) -> Dir:
     return Dir(path=dir_path, url=os.path.join(data_connection[0].gcs_folder.source, *dir_path.split("/")[4:]))
 
 
+def _resolve_lightning_storage(dir_path: str) -> Dir:
+    from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+    client = LightningClient(max_tries=2)
+
+    # Get the ids from env variables
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID", None)
+    if project_id is None:
+        raise RuntimeError("The `LIGHTNING_CLOUD_PROJECT_ID` couldn't be found from the environment variables.")
+
+    target_name = dir_path.split("/")[3]
+
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    data_connection = [dc for dc in data_connections if dc.name == target_name]
+
+    if not data_connection:
+        raise ValueError(f"We didn't find any matching data connection with the provided name `{target_name}`.")
+
+    return Dir(path=dir_path, url=os.path.join(data_connection[0].r2.source, *dir_path.split("/")[4:]))
+
+
 def _resolve_datasets(dir_path: str) -> Dir:
     from lightning_sdk.lightning_cloud.rest_client import LightningClient
 

tests/streaming/test_resolver.py

@@ -465,3 +465,45 @@ def test_src_resolver_gcs_folders(monkeypatch, lightning_cloud_mock):
     assert resolver._resolve_dir("/teamspace/gcs_folders/debug_folder/a/b/c").url == expected + "/a/b/c"
 
     auth.clear()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="windows isn't supported")
+def test_src_resolver_lightning_storage(monkeypatch, lightning_cloud_mock):
+    """Test lightning_storage resolver with r2 source."""
+    auth = login.Auth()
+    auth.save(user_id="7c8455e3-7c5f-4697-8a6d-105971d6b9bd", api_key="e63fae57-2b50-498b-bc46-d6204cbf330e")
+
+    with pytest.raises(
+        RuntimeError, match="`LIGHTNING_CLOUD_PROJECT_ID` couldn't be found from the environment variables."
+    ):
+        resolver._resolve_dir("/teamspace/lightning_storage/my_dataset")
+
+    monkeypatch.setenv("LIGHTNING_CLOUD_PROJECT_ID", "project_id")
+
+    client_mock = mock.MagicMock()
+    client_mock.data_connection_service_list_data_connections.return_value = V1ListDataConnectionsResponse(
+        data_connections=[V1DataConnection(name="my_dataset", r2=mock.MagicMock(source="r2://my-r2-bucket"))],
+    )
+
+    client_cls_mock = mock.MagicMock()
+    client_cls_mock.return_value = client_mock
+    lightning_cloud_mock.rest_client.LightningClient = client_cls_mock
+
+    expected = "r2://my-r2-bucket"
+    assert resolver._resolve_dir("/teamspace/lightning_storage/my_dataset").url == expected
+    assert resolver._resolve_dir("/teamspace/lightning_storage/my_dataset/train").url == expected + "/train"
+
+    # Test missing data connection
+    client_mock = mock.MagicMock()
+    client_mock.data_connection_service_list_data_connections.return_value = V1ListDataConnectionsResponse(
+        data_connections=[],
+    )
+
+    client_cls_mock = mock.MagicMock()
+    client_cls_mock.return_value = client_mock
+    lightning_cloud_mock.rest_client.LightningClient = client_cls_mock
+
+    with pytest.raises(ValueError, match="name `my_dataset`"):
+        resolver._resolve_dir("/teamspace/lightning_storage/my_dataset")
+
+    auth.clear()