ref: Move raw dataset code to litdata/raw, expose StreamingRawDataset at top-level (Lightning-AI#671)

* ref(constants): remove `asyncio` from requirement cache, already in the std lib

* ref: organize raw dataset files to separate folder

* ref: update imports and enhance __all__ exports in __init__.py

* ref: clarify comments for cloud and local file discovery in FileIndexer

* ref: move tests for raw dataet to `raw` folder with respective files

* refactor: update import statements for StreamingRawDataset and FileMetadata in usage examples

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add init files

* expose classes at module level

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bhimrazy

README.md

@@ -221,8 +221,8 @@ pip install "litdata[extra]" gcsfs
 
 **Usage Example:**
 ```python
-from litdata.streaming.raw_dataset import StreamingRawDataset
 from torch.utils.data import DataLoader
+from litdata import StreamingRawDataset
 
 dataset = StreamingRawDataset("s3://bucket/files/")
 
@@ -239,18 +239,19 @@ for batch in loader:
 You can also customize how files are grouped by subclassing `StreamingRawDataset` and overriding the `setup` method. This is useful for pairing related files (e.g., image and mask, audio and transcript) or any custom grouping logic.
 
 ```python
-from litdata.streaming.raw_dataset import StreamingRawDataset, FileMetadata
-from torch.utils.data import DataLoader
 from typing import Union
+from torch.utils.data import DataLoader
+from litdata import StreamingRawDataset
+from litdata.raw.indexer import FileMetadata
 
 class SegmentationRawDataset(StreamingRawDataset):
-  def setup(self, files: list[FileMetadata]) -> Union[list[FileMetadata], list[list[FileMetadata]]]:
-      # TODO: Implement your custom grouping logic here.
-      # For example, group files by prefix, extension, or any rule you need.
-      # Return a list of groups, where each group is a list of FileMetadata.
-      # Example:
-      #   return [[image, mask], ...]
-      pass
+    def setup(self, files: list[FileMetadata]) -> Union[list[FileMetadata], list[list[FileMetadata]]]:
+        # TODO: Implement your custom grouping logic here.
+        # For example, group files by prefix, extension, or any rule you need.
+        # Return a list of groups, where each group is a list of FileMetadata.
+        # Example:
+        #   return [[image, mask], ...]
+        pass
 
 # Initialize the custom dataset
 dataset = SegmentationRawDataset("s3://bucket/files/")

src/litdata/__init__.py

@@ -12,10 +12,10 @@
 # limitations under the License.
 import warnings
 
-from lightning_utilities.core.imports import RequirementCache
-
 from litdata.__about__ import *  # noqa: F403
+from litdata.constants import _LIGHTNING_SDK_AVAILABLE
 from litdata.processing.functions import map, merge_datasets, optimize, walk
+from litdata.raw.dataset import StreamingRawDataset
 from litdata.streaming.combined import CombinedStreamingDataset
 from litdata.streaming.dataloader import StreamingDataLoader
 from litdata.streaming.dataset import StreamingDataset
@@ -32,9 +32,9 @@
     category=UserWarning,
 )
 
-
 __all__ = [
     "StreamingDataset",
+    "StreamingRawDataset",
     "CombinedStreamingDataset",
     "StreamingDataLoader",
     "TokensLoader",
@@ -48,7 +48,8 @@
     "index_hf_dataset",
     "breakpoint",
 ]
-if RequirementCache("lightning_sdk"):
+
+if _LIGHTNING_SDK_AVAILABLE:
     from lightning_sdk import Machine  # noqa: F401
 
-    __all__ + ["Machine"]
+    __all__.append("Machine")

src/litdata/constants.py

@@ -45,7 +45,6 @@
 _TORCH_VISION_AVAILABLE = RequirementCache("torchvision")
 _AV_AVAILABLE = RequirementCache("av")
 _OBSTORE_AVAILABLE = RequirementCache("obstore")
-_ASYNCIO_AVAILABLE = RequirementCache("asyncio")
 
 _DEBUG = bool(int(os.getenv("DEBUG_LITDATA", "0")))
 _PRINT_DEBUG_LOGS = bool(int(os.getenv("PRINT_DEBUG_LOGS", "0")))

src/litdata/raw/__init__.py

@@ -0,0 +1,17 @@
+# Copyright The Lightning AI team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from litdata.raw.dataset import StreamingRawDataset
+from litdata.raw.indexer import FileMetadata
+
+__all__ = ["FileMetadata", "StreamingRawDataset"]

src/litdata/streaming/raw_dataset.py src/litdata/raw/dataset.pysrc/litdata/streaming/raw_dataset.py renamed to src/litdata/raw/dataset.py

@@ -10,185 +10,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import json
+
+import asyncio
 import logging
 import os
-import time
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
 from typing import Any, Callable, Optional, Union
-from urllib.parse import urlparse
 
 from torch.utils.data import Dataset
 
-from litdata.constants import _ASYNCIO_AVAILABLE, _FSSPEC_AVAILABLE, _TQDM_AVAILABLE, _ZSTD_AVAILABLE
+from litdata.raw.indexer import BaseIndexer, FileIndexer, FileMetadata
 from litdata.streaming.downloader import Downloader, get_downloader
 from litdata.streaming.resolver import Dir, _resolve_dir
 from litdata.utilities.dataset_utilities import generate_md5_hash, get_default_cache_dir
 
-if not _ASYNCIO_AVAILABLE:
-    raise ModuleNotFoundError(
-        "The 'asyncio' package is required for streaming datasets. Please install it with `pip install asyncio`."
-    )
-else:
-    import asyncio
-
 logger = logging.getLogger(__name__)
-SUPPORTED_PROVIDERS = ("s3", "gs", "azure")
-
-
-@dataclass
-class FileMetadata:
-    """Metadata for a single file in the dataset."""
-
-    path: str
-    size: int
-
-    def to_dict(self) -> dict[str, Any]:
-        return {"path": self.path, "size": self.size}
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "FileMetadata":
-        return cls(path=data["path"], size=data["size"])
-
-
-class BaseIndexer(ABC):
-    """Abstract base class for file indexing strategies."""
-
-    @abstractmethod
-    def discover_files(self, input_dir: str, storage_options: Optional[dict[str, Any]]) -> list[FileMetadata]:
-        """Discover dataset files and return their metadata."""
-
-    def build_or_load_index(
-        self, input_dir: str, cache_dir: str, storage_options: Optional[dict[str, Any]]
-    ) -> list[FileMetadata]:
-        """Build or load a ZSTD-compressed index of file metadata."""
-        if not _ZSTD_AVAILABLE:
-            raise ModuleNotFoundError(str(_ZSTD_AVAILABLE))
-
-        import zstd
-
-        index_path = Path(cache_dir) / "index.json.zstd"
-
-        # Try loading cached index if it exists
-        if index_path.exists():
-            try:
-                with open(index_path, "rb") as f:
-                    compressed_data = f.read()
-                metadata = json.loads(zstd.decompress(compressed_data).decode("utf-8"))
-
-                return [FileMetadata.from_dict(file_data) for file_data in metadata["files"]]
-            except (FileNotFoundError, json.JSONDecodeError, zstd.ZstdError, KeyError) as e:
-                logger.warning(f"Failed to load cached index from {index_path}: {e}")
-
-        # Build fresh index
-        logger.info(f"Building index for {input_dir} at {index_path}")
-        files = self.discover_files(input_dir, storage_options)
-        if not files:
-            raise ValueError(f"No files found in {input_dir}")
-
-        # Cache the index with ZSTD compression
-        # TODO: upload the index to cloud storage
-        try:
-            metadata = {
-                "source": input_dir,
-                "files": [file.to_dict() for file in files],
-                "created_at": time.time(),
-            }
-            with open(index_path, "wb") as f:
-                f.write(zstd.compress(json.dumps(metadata).encode("utf-8")))
-        except (OSError, zstd.ZstdError) as e:
-            logger.warning(f"Error caching index to {index_path}: {e}")
-
-        logger.info(f"Built index with {len(files)} files from {input_dir} at {index_path}")
-        return files
-
-
-class FileIndexer(BaseIndexer):
-    """Indexes files recursively from cloud or local storage with optional extension filtering."""
-
-    def __init__(
-        self,
-        max_depth: int = 5,
-        extensions: Optional[list[str]] = None,
-    ):
-        self.max_depth = max_depth
-        self.extensions = [ext.lower() for ext in (extensions or [])]
-
-    def discover_files(self, input_dir: str, storage_options: Optional[dict[str, Any]]) -> list[FileMetadata]:
-        """Discover dataset files and return their metadata."""
-        parsed_url = urlparse(input_dir)
-
-        if parsed_url.scheme in SUPPORTED_PROVIDERS:
-            return self._discover_cloud_files(input_dir, storage_options)
-
-        if not parsed_url.scheme or parsed_url.scheme == "file":
-            return self._discover_local_files(input_dir)
-
-        raise ValueError(
-            f"Unsupported input directory scheme: {parsed_url.scheme}. Supported schemes are: {SUPPORTED_PROVIDERS}"
-        )
-
-    def _discover_cloud_files(self, input_dir: str, storage_options: Optional[dict[str, Any]]) -> list[FileMetadata]:
-        """Recursively list files in a cloud storage bucket."""
-        if not _FSSPEC_AVAILABLE:
-            raise ModuleNotFoundError(str(_FSSPEC_AVAILABLE))
-        import fsspec
-
-        obj = urlparse(input_dir)
-
-        # TODO: Research on switching to 'obstore' for file listing to potentially improve performance.
-        # Currently using 'fsspec' due to some issues with 'obstore' when handling multiple instances.
-        fs = fsspec.filesystem(obj.scheme, **(storage_options or {}))
-        files = fs.find(input_dir, maxdepth=self.max_depth, detail=True, withdirs=False)
-
-        if _TQDM_AVAILABLE:
-            from tqdm.auto import tqdm
-
-            pbar = tqdm(desc="Discovering files", total=len(files))
-
-        metadatas = []
-        for _, file_info in files.items():
-            if file_info.get("type") != "file":
-                continue
-
-            file_path = file_info["name"]
-            if self._should_include_file(file_path):
-                metadata = FileMetadata(
-                    path=f"{obj.scheme}://{file_path}",
-                    size=file_info.get("size", 0),
-                )
-                metadatas.append(metadata)
-            if _TQDM_AVAILABLE:
-                pbar.update(1)
-        if _TQDM_AVAILABLE:
-            pbar.close()
-        return metadatas
-
-    def _discover_local_files(self, input_dir: str) -> list[FileMetadata]:
-        """Recursively list files in the local filesystem."""
-        path = Path(input_dir)
-        metadatas = []
-
-        for file_path in path.rglob("*"):
-            if not file_path.is_file():
-                continue
-
-            if self._should_include_file(str(file_path)):
-                metadata = FileMetadata(
-                    path=str(file_path),
-                    size=file_path.stat().st_size,
-                )
-                metadatas.append(metadata)
-
-        return metadatas
-
-    def _should_include_file(self, file_path: str) -> bool:
-        """Return True if file matches allowed extensions."""
-        file_ext = Path(file_path).suffix.lower()
-        return not self.extensions or file_ext in self.extensions
 
 
 class CacheManager: