Data Processor: Add is_last argument to know when the last item for the current worker is being processed (#19383)

Author: tchaton

src/lightning/data/streaming/data_processor.py

@@ -608,7 +608,7 @@ def _handle_data_chunk_recipe_end(self) -> None:
     def _handle_data_transform_recipe(self, index: int) -> None:
         # Don't use a context manager to avoid deleting files that are being uploaded.
         output_dir = tempfile.mkdtemp()
-        item_data = self.data_recipe.prepare_item(self.items[index], str(output_dir))
+        item_data = self.data_recipe.prepare_item(self.items[index], str(output_dir), len(self.items) - 1 == index)
         if item_data is not None:
             raise ValueError(
                 "When using a `DataTransformRecipe`, the `prepare_item` shouldn't return anything."
@@ -649,33 +649,9 @@ def prepare_structure(self, input_dir: Optional[str]) -> List[T]:
         pass
 
     @abstractmethod
-    def prepare_item(self, *args: Any) -> Any:
+    def prepare_item(self, *args: Any, **kwargs: Any) -> Any:
         pass
 
-    def listdir(self, path: str) -> List[str]:
-        home = _get_home_folder()
-        filepath = os.path.join(home, ".cache", f"{self._name}/filepaths.txt")
-
-        if os.path.exists(filepath):
-            lines = []
-            with open(filepath) as f:
-                for line in f.readlines():
-                    lines.append(line.replace("\n", ""))
-            return lines
-
-        filepaths = []
-        for dirpath, _, filenames in os.walk(path):
-            for filename in filenames:
-                filepaths.append(os.path.join(dirpath, filename))
-
-        os.makedirs(os.path.dirname(filepath), exist_ok=True)
-
-        with open(filepath, "w") as f:
-            for filepath in filepaths:
-                f.write(f"{filepath}\n")
-
-        return filepaths
-
     def __init__(self) -> None:
         self._name: Optional[str] = None
 
@@ -707,7 +683,7 @@ def prepare_structure(self, input_dir: Optional[str]) -> List[T]:
         """
 
     @abstractmethod
-    def prepare_item(self, item_metadata: T) -> Any:  # type: ignore
+    def prepare_item(self, item_metadata: T) -> Any:
         """The return of this `prepare_item` method is persisted in chunked binary files."""
 
     def _done(self, size: int, delete_cached_files: bool, output_dir: Dir) -> _Result:
@@ -798,7 +774,7 @@ def prepare_structure(self, input_dir: Optional[str]) -> List[T]:
         """
 
     @abstractmethod
-    def prepare_item(self, item_metadata: T, output_dir: str) -> None:  # type: ignore
+    def prepare_item(self, item_metadata: T, output_dir: str, is_last: bool) -> None:
         """Use your item metadata to process your files and save the file outputs into `output_dir`."""
 
 

src/lightning/data/streaming/functions.py

@@ -78,25 +78,28 @@ def __init__(self, fn: Callable[[str, Any], None], inputs: Sequence[Any]):
         _fn = self._fn if isinstance(self._fn, FunctionType) else self._fn.__call__  # type: ignore
         params = inspect.signature(_fn).parameters
         self._contains_device = "device" in params
+        self._contains_is_last = "is_last" in params
 
     def prepare_structure(self, _: Optional[str]) -> Any:
         return self._inputs
 
-    def prepare_item(self, item_metadata: Any, output_dir: str) -> None:  # type: ignore
+    def prepare_item(self, item_metadata: Any, output_dir: str, is_last: bool) -> None:
         if self._contains_device and self._device is None:
             self._find_device()
 
+        kwargs: Dict[str, Any] = {}
+
+        if self._contains_device:
+            kwargs["device"] = self._device
+
+        if self._contains_is_last:
+            kwargs["is_last"] = is_last
+
         if isinstance(self._fn, (FunctionType, partial)):
-            if self._contains_device:
-                self._fn(item_metadata, output_dir, self._device)
-            else:
-                self._fn(item_metadata, output_dir)
+            self._fn(item_metadata, output_dir, **kwargs)
 
         elif callable(self._fn):
-            if self._contains_device:
-                self._fn.__call__(item_metadata, output_dir, self._device)  # type: ignore
-            else:
-                self._fn.__call__(item_metadata, output_dir)  # type: ignore
+            self._fn.__call__(item_metadata, output_dir, **kwargs)  # type: ignore
         else:
             raise ValueError(f"The provided {self._fn} isn't supported.")
 
@@ -124,7 +127,7 @@ def __init__(
     def prepare_structure(self, input_dir: Optional[str]) -> Any:
         return self._inputs
 
-    def prepare_item(self, item_metadata: Any) -> Any:  # type: ignore
+    def prepare_item(self, item_metadata: Any) -> Any:
         if isinstance(self._fn, partial):
             yield from self._fn(item_metadata)
 

tests/tests_data/streaming/test_data_processor.py

@@ -350,7 +350,7 @@ def test_map_items_to_workers_sequentially(monkeypatch):
 
 class CustomDataChunkRecipe(DataChunkRecipe):
     def prepare_structure(self, input_dir: str) -> List[Any]:
-        filepaths = self.listdir(input_dir)
+        filepaths = [os.path.join(input_dir, f) for f in os.listdir(input_dir)]
         assert len(filepaths) == 30
         return filepaths
 
@@ -567,7 +567,7 @@ def prepare_structure(self, input_dir: str):
         filepaths = [os.path.join(input_dir, filename) for filename in os.listdir(input_dir)]
         return [filepath for filepath in filepaths if os.path.isfile(filepath)]
 
-    def prepare_item(self, filepath: Any, output_dir: str) -> None:
+    def prepare_item(self, filepath: Any, output_dir: str, is_last) -> None:
         from PIL import Image
 
         img = Image.open(filepath)
@@ -819,7 +819,7 @@ def fn(output_dir, item, device):
 
     data_recipe = LambdaDataTransformRecipe(fn, range(1))
 
-    data_recipe.prepare_item(1, "")
+    data_recipe.prepare_item(1, "", False)
     assert called
 
 
@@ -839,7 +839,7 @@ def __call__(self, item, output_dir, device):
             called = True
 
     data_recipe = LambdaDataTransformRecipe(Transform(), range(1))
-    data_recipe.prepare_item(1, "")
+    data_recipe.prepare_item(1, "", False)
     assert called
 
 
@@ -968,7 +968,7 @@ def test_data_processing_map_non_absolute_path(monkeypatch, tmpdir):
 
 
 @pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows")
-def test_map_error_when_not_empty(monkeypatch, tmpdir):
+def test_map_error_when_not_empty(monkeypatch):
     boto3 = mock.MagicMock()
     client_s3_mock = mock.MagicMock()
     client_s3_mock.list_objects_v2.return_value = {"KeyCount": 1, "Contents": []}
@@ -992,3 +992,27 @@ def test_map_error_when_not_empty(monkeypatch, tmpdir):
             output_dir=Dir(path=None, url="s3://bucket"),
             error_when_not_empty=False,
         )
+
+def map_fn_is_last(index, output_dir, is_last):
+    with open(os.path.join(output_dir, f"{index}_{is_last}.txt"), "w") as f:
+        f.write("here")
+
+
+@pytest.mark.skipif(condition=sys.platform == "win32", reason="Not supported on windows")
+@pytest.mark.parametrize(
+    ("num_workers", "expected"),
+    [
+        (1, ['0_False.txt', '1_False.txt', '2_False.txt', '3_False.txt', '4_True.txt']),
+        (2, ['0_False.txt', '1_True.txt', '2_False.txt', '3_False.txt', '4_True.txt']),
+    ],
+)
+def test_map_is_last(num_workers, expected, tmpdir):
+    map(
+        map_fn_is_last,
+        list(range(5)),
+        output_dir=str(tmpdir),
+        error_when_not_empty=False,
+        num_workers=num_workers,
+    )
+
+    assert sorted(os.listdir(tmpdir)) == expected