update

Lightning-AI · Feb 23, 2024 · fcab3bd · fcab3bd
1 parent f2f3ef5
commit fcab3bd
Show file tree

Hide file tree

Showing 55 changed files with 13 additions and 11,264 deletions.
diff --git a/.github/workflows/ci-tests-data.yml b/.github/workflows/ci-tests-data.yml
diff --git a/requirements/data/data.txt b/requirements/data/data.txt
@@ -6,3 +6,4 @@ lightning-utilities >=0.8.0, <0.10.0
 torch >0.14.0, <=2.2.0
 lightning-cloud
 filelock
+litdata >= 0.2.0rc
diff --git a/src/lightning/data/__init__.py b/src/lightning/data/__init__.py
@@ -2,26 +2,20 @@
 
 from lightning_utilities.core.imports import RequirementCache
 
-_LIGHTNING_DATA_AVAILABLE = RequirementCache("lightning_data")
-_LIGHTNING_SDK_AVAILABLE = RequirementCache("lightning_sdk")
+_LIDATA_AVAILABLE = RequirementCache("litdata")
 
-if _LIGHTNING_DATA_AVAILABLE:
-    import lightning_data
+if not _LIDATA_AVAILABLE:
+    raise ModuleNotFoundError("Please, run `pip install litdata`")  # E111
 
-    # Enable resolution at least for lower data namespace
-    sys.modules["lightning.data"] = lightning_data
+import litdata  # noqa: E402
 
-    from lightning_data.processing.functions import map, optimize, walk
-    from lightning_data.streaming.combined import CombinedStreamingDataset
-    from lightning_data.streaming.dataloader import StreamingDataLoader
-    from lightning_data.streaming.dataset import StreamingDataset
+# Enable resolution at least for lower data namespace
+sys.modules["lightning.data"] = litdata
 
-else:
-    # TODO: Delete all the code when everything is moved to lightning_data
-    from lightning.data.processing.functions import map, optimize, walk
-    from lightning.data.streaming.combined import CombinedStreamingDataset
-    from lightning.data.streaming.dataloader import StreamingDataLoader
-    from lightning.data.streaming.dataset import StreamingDataset
+from litdata.processing.functions import map, optimize, walk  # noqa: E402
+from litdata.streaming.combined import CombinedStreamingDataset  # noqa: E402
+from litdata.streaming.dataloader import StreamingDataLoader  # noqa: E402
+from litdata.streaming.dataset import StreamingDataset  # noqa: E402
 
 __all__ = [
     "LightningDataset",
@@ -34,8 +28,8 @@
     "walk",
 ]
 
-# TODO: Move this to lightning_data
-if _LIGHTNING_SDK_AVAILABLE:
+# TODO: Move this to litdata
+if RequirementCache("lightning_sdk"):
     from lightning_sdk import Machine  # noqa: F401
 
     __all__.append("Machine")
diff --git a/src/lightning/data/constants.py b/src/lightning/data/constants.py
diff --git a/src/lightning/data/processing/__init__.py b/src/lightning/data/processing/__init__.py