Merge branch 'master' into bugfix/csv-file-exists

Author: awaelchli

requirements/data/data.txt

@@ -3,6 +3,6 @@
 
 lightning-utilities >=0.8.0, <0.10.0
 # to be able to include also PL 2.0 and preserve `>` needed for CI min version bypass
-torch >0.14.0, <2.2.0
+torch >0.14.0, <=2.2.0
 lightning-cloud
 filelock

requirements/data/examples.txt

@@ -1,3 +1,3 @@
 Pillow >= 9.5.0
 # min version to match torch >= 2.0.1
-torchvision >=0.15.2, <0.17.0
+torchvision >=0.15.2, <0.18.0

requirements/data/test.txt

@@ -5,5 +5,5 @@ pytest-timeout ==2.1.0
 pytest-rerunfailures ==12.0
 pytest-random-order ==1.1.0
 viztracer
+pandas
 pyarrow
-polars

requirements/fabric/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <2.2.0
+torch >=1.13.0, <=2.2.0
 fsspec[http] >=2022.5.0, <2023.11.0
 packaging >=20.0, <=23.1
 typing-extensions >=4.4.0, <4.10.0

requirements/fabric/examples.txt

@@ -1,6 +1,6 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torchvision >=0.14.0, <0.17.0
+torchvision >=0.14.0, <0.18.0
 torchmetrics >=0.10.0, <1.3.0
 lightning-utilities >=0.8.0, <0.10.0

requirements/pytorch/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <2.2.0
+torch >=1.13.0, <=2.2.0
 tqdm >=4.57.0, <4.67.0
 PyYAML >=5.4, <6.1.0
 fsspec[http] >=2022.5.0, <2023.11.0

requirements/pytorch/examples.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 requests <2.32.0
-torchvision >=0.14.0, <0.17.0
+torchvision >=0.14.0, <0.18.0
 gym[classic_control] >=0.17.0, <0.27.0
 ipython[all] <8.15.0
 torchmetrics >=0.10.0, <1.3.0

requirements/typing.txt

@@ -1,5 +1,5 @@
-mypy==1.5.1
-torch==2.1.0
+mypy==1.8.0
+torch==2.2.0
 
 types-Markdown
 types-PyYAML

src/lightning/app/cli/lightning_cli.py

@@ -363,7 +363,7 @@ def init_pl_app(source: Union[Tuple[str], Tuple[str, str]], name: str, overwrite
         source_dir = str(Path(script_path).resolve().parent)
     elif len(source) == 2:
         # enable type checking once https://github.com/python/mypy/issues/1178 is available
-        source_dir, script_path = source  # type: ignore
+        source_dir, script_path = source
     else:
         click.echo(
             f"Incorrect number of arguments. You passed ({', '.join(source)}) but only either one argument"

src/lightning/data/__init__.py

@@ -1,3 +1,5 @@
+from lightning_utilities.core.imports import RequirementCache
+
 from lightning.data.processing.functions import map, optimize, walk
 from lightning.data.streaming.combined import CombinedStreamingDataset
 from lightning.data.streaming.dataloader import StreamingDataLoader
@@ -13,3 +15,8 @@
     "optimize",
     "walk",
 ]
+
+if RequirementCache('lightning_sdk'):
+    from lightning_sdk import Machine  # noqa: F401
+
+    __all__.append("Machine")

src/lightning/data/processing/data_processor.py

@@ -372,7 +372,6 @@ def __init__(
         self._counter = 0
         self._last_time = time()
         self._index_counter = 0
-        self._current_item: Any = None
 
     def run(self) -> None:
         try:
@@ -477,6 +476,7 @@ def _try_upload(self, data: Optional[Union[str, Tuple[str, str]]]) -> None:
             assert os.path.exists(data), data
         else:
             assert os.path.exists(data[-1]), data
+
         self.to_upload_queues[self._counter % self.num_uploaders].put(data)
 
     def _collect_paths(self) -> None:
@@ -588,8 +588,8 @@ def _start_uploaders(self) -> None:
 
     def _handle_data_chunk_recipe(self, index: int) -> None:
         try:
-            self._current_item = self.items[index] if self.reader is None else self.reader.read(self.items[index])
-            item_data_or_generator = self.data_recipe.prepare_item(self._current_item)
+            current_item = self.items[index] if self.reader is None else self.reader.read(self.items[index])
+            item_data_or_generator = self.data_recipe.prepare_item(current_item)
             if isinstance(item_data_or_generator, types.GeneratorType):
                 for item_data in item_data_or_generator:
                     if item_data is not None:
@@ -713,14 +713,19 @@ def _done(self, size: int, delete_cached_files: bool, output_dir: Dir) -> _Resul
             size = sum([c["dim"] if c["dim"] is not None else c["chunk_size"] for c in config["chunks"]])
             num_bytes = sum([c["chunk_bytes"] for c in config["chunks"]])
             data_format = tree_unflatten(config["config"]["data_format"], treespec_loads(config["config"]["data_spec"]))
+            num_chunks = len(config["chunks"])
+
+            # The platform can't store more than 1024 entries.
+            # Note: This isn't really used right now, so it is fine to skip if too big.
+            num_bytes_per_chunk = [c["chunk_size"] for c in config["chunks"]] if num_chunks < 1024 else []
 
             return _Result(
                 size=size,
                 num_bytes=num_bytes,
                 data_format=data_format,
                 compression=config["config"]["compression"],
                 num_chunks=len(config["chunks"]),
-                num_bytes_per_chunk=[c["chunk_size"] for c in config["chunks"]],
+                num_bytes_per_chunk=num_bytes_per_chunk,
             )
         return _Result(
             size=size,
@@ -866,9 +871,9 @@ def run(self, data_recipe: DataRecipe) -> None:
             raise ValueError("The `prepare_structure` should return a list of item metadata.")
 
         if self.reader:
-            workers_user_items = self.reader.items_to_workers(user_items, self.num_workers)
+            user_items = self.reader.remap_items(user_items, self.num_workers)
 
-        elif self.weights is not None:
+        if self.weights is not None:
             if len(self.weights) != len(user_items):
                 raise ValueError("The provided weights length should match the inputs' length.")
             workers_user_items = _map_items_to_workers_weighted(

src/lightning/data/processing/dns.py

@@ -24,8 +24,8 @@
 
 from lightning.data.constants import _IS_IN_STUDIO, _TORCH_GREATER_EQUAL_2_1_0
 from lightning.data.processing.data_processor import DataChunkRecipe, DataProcessor, DataTransformRecipe
-from lightning.data.processing.dns import optimize_dns_context
 from lightning.data.processing.readers import BaseReader
+from lightning.data.processing.utilities import optimize_dns_context
 from lightning.data.streaming.resolver import (
     Dir,
     _assert_dir_has_index_file,