Streamingdataset torch compatibility (#108)

Author: yhl48

.github/workflows/ci-checks.yml

@@ -30,7 +30,7 @@ jobs:
       artifact-name: dist-packages-${{ github.sha }}
       testing-matrix: |
         {
-          "os": ["ubuntu-latest", "macos-latest", "windows-latest"],
+          "os": ["ubuntu-latest", "macos-13", "windows-latest"],
           "python-version": ["3.8", "3.10"]
         }
 

.github/workflows/ci-testing.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, macOS-latest, windows-latest]
+        os: [ubuntu-latest, macos-13, windows-latest]
         python-version: [3.9]
         requires: ["oldest", "latest"]
 

pyproject.toml

@@ -72,7 +72,8 @@ lint.ignore = [
 exclude = [
     ".git",
     "docs",
-    "_notebooks"
+    "_notebooks",
+    "src/litdata/utilities/_pytree.py",
 ]
 lint.ignore-init-module-imports = true
 

src/litdata/processing/data_processor.py

@@ -39,7 +39,6 @@
     _INDEX_FILENAME,
     _IS_IN_STUDIO,
     _LIGHTNING_CLOUD_AVAILABLE,
-    _TORCH_GREATER_EQUAL_2_1_0,
 )
 from litdata.imports import RequirementCache
 from litdata.processing.readers import BaseReader, StreamingDataLoaderReader
@@ -49,6 +48,7 @@
 from litdata.streaming.client import S3Client
 from litdata.streaming.dataloader import StreamingDataLoader
 from litdata.streaming.resolver import _resolve_dir
+from litdata.utilities._pytree import tree_flatten, tree_unflatten, treespec_loads
 from litdata.utilities.broadcast import broadcast_object
 from litdata.utilities.packing import _pack_greedily
 
@@ -57,9 +57,6 @@
 if _TQDM_AVAILABLE:
     from tqdm.auto import tqdm as _tqdm
 
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import tree_flatten, tree_unflatten, treespec_loads
-
 if _LIGHTNING_CLOUD_AVAILABLE:
     from lightning_cloud.openapi import V1DatasetType
 

src/litdata/processing/functions.py

@@ -22,7 +22,7 @@
 
 import torch
 
-from litdata.constants import _IS_IN_STUDIO, _TORCH_GREATER_EQUAL_2_1_0
+from litdata.constants import _IS_IN_STUDIO
 from litdata.processing.data_processor import DataChunkRecipe, DataProcessor, DataTransformRecipe
 from litdata.processing.readers import BaseReader
 from litdata.processing.utilities import optimize_dns_context
@@ -34,9 +34,7 @@
     _execute,
     _resolve_dir,
 )
-
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import tree_flatten
+from litdata.utilities._pytree import tree_flatten
 
 
 def _get_indexed_paths(data: Any) -> Dict[int, str]:

src/litdata/streaming/cache.py

@@ -17,7 +17,6 @@
 
 from litdata.constants import (
     _INDEX_FILENAME,
-    _TORCH_GREATER_EQUAL_2_1_0,
 )
 from litdata.streaming.item_loader import BaseItemLoader
 from litdata.streaming.reader import BinaryReader
@@ -56,9 +55,6 @@ def __init__(
 
         """
         super().__init__()
-        if not _TORCH_GREATER_EQUAL_2_1_0:
-            raise ModuleNotFoundError("PyTorch version 2.1 or higher is required to use the cache.")
-
         input_dir = _resolve_dir(input_dir)
         self._cache_dir = input_dir.path
         assert self._cache_dir

src/litdata/streaming/config.py

@@ -15,15 +15,13 @@
 import os
 from typing import Any, Dict, List, Optional, Tuple
 
-from litdata.constants import _INDEX_FILENAME, _TORCH_GREATER_EQUAL_2_1_0
+from litdata.constants import _INDEX_FILENAME
 from litdata.streaming.compression import _COMPRESSORS, Compressor
 from litdata.streaming.downloader import get_downloader_cls
 from litdata.streaming.item_loader import BaseItemLoader, PyTreeLoader, TokensLoader
 from litdata.streaming.sampler import ChunkedIndex
 from litdata.streaming.serializers import Serializer
-
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import tree_unflatten, treespec_loads
+from litdata.utilities._pytree import tree_unflatten, treespec_loads
 
 
 class ChunksConfig:

src/litdata/streaming/dataloader.py

@@ -33,7 +33,7 @@
 )
 from torch.utils.data.sampler import BatchSampler, Sampler
 
-from litdata.constants import _DEFAULT_CHUNK_BYTES, _TORCH_GREATER_EQUAL_2_1_0, _VIZ_TRACKER_AVAILABLE
+from litdata.constants import _DEFAULT_CHUNK_BYTES, _VIZ_TRACKER_AVAILABLE
 from litdata.streaming import Cache
 from litdata.streaming.combined import (
     __NUM_SAMPLES_YIELDED_KEY__,
@@ -42,11 +42,9 @@
 )
 from litdata.streaming.dataset import StreamingDataset
 from litdata.streaming.sampler import CacheBatchSampler
+from litdata.utilities._pytree import tree_flatten
 from litdata.utilities.env import _DistributedEnv
 
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import tree_flatten
-
 logger = logging.Logger(__name__)
 
 

src/litdata/streaming/item_loader.py

@@ -22,12 +22,9 @@
 
 from litdata.constants import (
     _TORCH_DTYPES_MAPPING,
-    _TORCH_GREATER_EQUAL_2_1_0,
 )
 from litdata.streaming.serializers import Serializer
-
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import PyTree, tree_unflatten
+from litdata.utilities._pytree import PyTree, tree_unflatten
 
 
 class BaseItemLoader(ABC):

src/litdata/streaming/writer.py

@@ -21,16 +21,14 @@
 import numpy as np
 import torch
 
-from litdata.constants import _INDEX_FILENAME, _TORCH_GREATER_EQUAL_2_1_0
+from litdata.constants import _INDEX_FILENAME
 from litdata.processing.utilities import get_worker_rank
 from litdata.streaming.compression import _COMPRESSORS, Compressor
 from litdata.streaming.serializers import Serializer, _get_serializers
+from litdata.utilities._pytree import PyTree, tree_flatten, treespec_dumps
 from litdata.utilities.env import _DistributedEnv, _WorkerEnv
 from litdata.utilities.format import _convert_bytes_to_int, _human_readable_bytes
 
-if _TORCH_GREATER_EQUAL_2_1_0:
-    from torch.utils._pytree import PyTree, tree_flatten, treespec_dumps
-
 
 @dataclass
 class Item: