Update Prefetcher and Implement PinMemory IterDataPipe (#1014)

Summary:
Fixes #1013

## Changes

- Simplify the control flow of prefetcher
  - Delay Exception raised from thread worker to main thread in `__iter__`
  - Stop prefetching whenever Exception is received
  - As long as `stop_iteration` is not turned on or `buffer` is not empty, continue yielding data from `__iter__`.
  - Add serialization test
- Add `PinMemory` DataPipe
  -  `is_replciable() -> False` to keep it in the main process
  - Add unit tests
- Update `test_proto_multi_rs.py` to `test_mprs.py`

Pull Request resolved: #1014

Reviewed By: NivekT

Differential Revision: D43329696

Pulled By: ejguan

fbshipit-source-id: da4326dbe2388f4e23b9a1a3a5c43da09d29185a

Author: ejguan

docs/source/index.rst

@@ -42,6 +42,7 @@ Features described in this documentation are classified by release status:
    dataloader2.rst
    reading_service.rst
 
+
 .. toctree::
    :maxdepth: 2
    :caption: Tutorial and Examples:

docs/source/torchdata.datapipes.utils.rst

@@ -15,6 +15,17 @@ DataPipe Graph Visualization
 
     to_graph
 
+Commond Utility Functions
+--------------------------------------
+.. currentmodule:: torchdata.datapipes.utils
+
+.. autosummary::
+    :nosignatures:
+    :toctree: generated/
+    :template: function.rst
+
+    pin_memory_fn
+
 
 File Object and Stream Utility
 -------------------------------------

test/dataloader2/test_proto_multi_rs.py renamed to test/dataloader2/test_mprs.py

@@ -10,7 +10,7 @@
 from unittest import TestCase
 
 from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize
-from torchdata.dataloader2 import DataLoader2, DataLoader2Iterator, PrototypeMultiProcessingReadingService
+from torchdata.dataloader2 import DataLoader2, DataLoader2Iterator, MultiProcessingReadingService
 from torchdata.datapipes.iter import IterableWrapper
 
 
@@ -29,9 +29,9 @@ def _add_one(x: int) -> int:
 dp_parametrize = parametrize("dp", test_dps)
 
 
-class TestPrototypeMultiProcessingReadingService(TestCase):
+class TestMultiProcessingReadingService(TestCase):
     r"""
-    This tests specific functionalities of PrototypeMultiProcessingReadingService, notably
+    This tests specific functionalities of MultiProcessingReadingService, notably
     `pause`, `resume`, `snapshot`.
     """
 
@@ -40,7 +40,7 @@ def test_reading_service_pause_resume_0_worker(self, ctx) -> None:
 
         # Functional Test: Verifies that this ReadingService will raise error when `pause/resume` is used
         #                  with `num_workers = 0`
-        rs0 = PrototypeMultiProcessingReadingService(
+        rs0 = MultiProcessingReadingService(
             num_workers=0, worker_prefetch_cnt=0, main_prefetch_cnt=0, multiprocessing_context=ctx
         )
         dl0: DataLoader2 = DataLoader2(dp1, reading_service=rs0)
@@ -64,7 +64,7 @@ def test_reading_service_pause_resume(self, ctx, dp, n_workers, worker_prefetch_
 
         # Functional Test: Testing various configuration of DataPipe/ReadingService to ensure the pipeline
         #                  properly pauses and resumes
-        rs = PrototypeMultiProcessingReadingService(
+        rs = MultiProcessingReadingService(
             num_workers=n_workers,
             worker_prefetch_cnt=worker_prefetch_cnt,
             main_prefetch_cnt=main_prefetch_cnt,
@@ -93,7 +93,7 @@ def test_reading_service_pause_resume(self, ctx, dp, n_workers, worker_prefetch_
     def test_reading_service_pause_stop_yield(self, ctx, dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None:
 
         # Functional Test: Confirms that `dl` will stop yielding elements after `_pause` is called
-        rs = PrototypeMultiProcessingReadingService(
+        rs = MultiProcessingReadingService(
             num_workers=n_workers,
             worker_prefetch_cnt=worker_prefetch_cnt,
             main_prefetch_cnt=main_prefetch_cnt,
@@ -117,7 +117,7 @@ def test_reading_service_pause_stop_yield(self, ctx, dp, n_workers, worker_prefe
     @parametrize("n_workers,worker_prefetch_cnt,main_prefetch_cnt", [(1, 0, 0), (1, 0, 2), (2, 0, 0), (2, 2, 2)])
     def test_reading_service_limit(self, dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None:
 
-        rs = PrototypeMultiProcessingReadingService(
+        rs = MultiProcessingReadingService(
             num_workers=n_workers, worker_prefetch_cnt=worker_prefetch_cnt, main_prefetch_cnt=main_prefetch_cnt
         )
 
@@ -209,10 +209,10 @@ def test_reading_service_limit(self, dp, n_workers, worker_prefetch_cnt, main_pr
     #       those DPs belong to a dispatching process and only do pause if worker_id == 0
     #       There might still be a race condition, need to look into the messages
 
-    # rs1 = PrototypeMultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=0, main_prefetch_cnt=0)
-    # rs2 = PrototypeMultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=0, main_prefetch_cnt=2)
-    # rs3 = PrototypeMultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=2, main_prefetch_cnt=0)
-    # rs4 = PrototypeMultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=2, main_prefetch_cnt=2)
+    # rs1 = MultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=0, main_prefetch_cnt=0)
+    # rs2 = MultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=0, main_prefetch_cnt=2)
+    # rs3 = MultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=2, main_prefetch_cnt=0)
+    # rs4 = MultiProcessingReadingService(num_workers=2, worker_prefetch_cnt=2, main_prefetch_cnt=2)
     # rss = [rs1, rs2, rs3, rs4]
 
     # for n, rs in enumerate(rss):
@@ -284,7 +284,7 @@ def test_reading_service_limit(self, dp, n_workers, worker_prefetch_cnt, main_pr
     #     pass
 
 
-instantiate_parametrized_tests(TestPrototypeMultiProcessingReadingService)
+instantiate_parametrized_tests(TestMultiProcessingReadingService)
 
 
 if __name__ == "__main__":

test/test_iterdatapipe.py

@@ -14,7 +14,7 @@
 from typing import Dict
 
 import expecttest
-import torch.utils.data.datapipes.iter
+import torch
 
 import torchdata
 
@@ -42,6 +42,8 @@
 )
 from torchdata.datapipes.map import MapDataPipe, SequenceWrapper
 
+skipIfNoCUDA = unittest.skipIf(not torch.cuda.is_available(), "CUDA is not available")
+
 
 def test_torchdata_pytorch_consistency() -> None:
     def extract_datapipe_names(module):
@@ -68,6 +70,14 @@ def extract_datapipe_names(module):
         raise AssertionError(msg + "\n".join(sorted(missing_datapipes)))
 
 
+def _convert_to_tensor(data):
+    if isinstance(data, dict):
+        return {k: _convert_to_tensor(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [_convert_to_tensor(v) for v in data]
+    return torch.tensor(data)
+
+
 class TestIterDataPipe(expecttest.TestCase):
     def test_in_memory_cache_holder_iterdatapipe(self) -> None:
         source_dp = IterableWrapper(range(10))
@@ -1475,6 +1485,38 @@ def test_random_splitter_iterdatapipe(self):
             next(it_train)
         next(it_valid)  # No error, can keep going
 
+    @skipIfNoCUDA
+    def test_pin_memory(self):
+        # Tensor
+        dp = IterableWrapper([(i, i + 1) for i in range(10)]).map(_convert_to_tensor).pin_memory()
+        self.assertTrue(all(d.is_pinned() for d in dp))
+
+        # List of Tensors
+        dp = IterableWrapper([[(i - 1, i), (i, i + 1)] for i in range(10)]).map(_convert_to_tensor).pin_memory()
+        self.assertTrue(all(d0.is_pinned() and d1.is_pinned() for d0, d1 in dp))
+
+        # Dict of Tensors
+        dp = IterableWrapper([{str(i): (i, i + 1)} for i in range(10)]).map(_convert_to_tensor).pin_memory()
+        self.assertTrue(all(v.is_pinned() for d in dp for v in d.values()))
+
+        # Dict of List of Tensors
+        dp = (
+            IterableWrapper([{str(i): [(i - 1, i), (i, i + 1)]} for i in range(10)])
+            .map(_convert_to_tensor)
+            .pin_memory()
+        )
+        self.assertTrue(all(v.is_pinned() for d in dp for batch in d.values() for v in batch))
+
+        # List of Dict of Tensors
+        dp = IterableWrapper([{str(i): (i, i + 1)} for i in range(10)]).map(_convert_to_tensor).batch(2).pin_memory()
+        self.assertTrue(all(v.is_pinned() for batch in dp for d in batch for v in d.values()))
+
+        # List of List of Tensors
+        dp = (
+            IterableWrapper([[(i - 1, i), (i, i + 1)] for i in range(10)]).map(_convert_to_tensor).batch(2).pin_memory()
+        )
+        self.assertTrue(all(d0.is_pinned() and d1.is_pinned() for batch in dp for d0, d1 in batch))
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_serialization.py

@@ -92,6 +92,10 @@ def _filter_by_module_availability(datapipes):
     return [dp for dp in datapipes if dp[0] not in filter_set]
 
 
+def _convert_to_tensor(data):
+    return torch.tensor(data)
+
+
 class TestIterDataPipeSerialization(expecttest.TestCase):
     def setUp(self):
         self.temp_dir = create_temp_dir()
@@ -272,6 +276,7 @@ def test_serializable(self):
                 (),
                 {},
             ),
+            (iterdp.Prefetcher, None, (), {}),
             (iterdp.ParquetDataFrameLoader, None, (), {"dtype": DTYPE}),
             (iterdp.RarArchiveLoader, None, (), {}),
             (

torchdata/datapipes/iter/__init__.py

@@ -108,7 +108,10 @@
     CSVParserIterDataPipe as CSVParser,
     LineReaderIterDataPipe as LineReader,
 )
-from torchdata.datapipes.iter.util.prefetcher import PrefetcherIterDataPipe as Prefetcher
+from torchdata.datapipes.iter.util.prefetcher import (
+    PinMemoryIterDataPipe as PinMemory,
+    PrefetcherIterDataPipe as Prefetcher,
+)
 from torchdata.datapipes.iter.util.randomsplitter import RandomSplitterIterDataPipe as RandomSplitter
 from torchdata.datapipes.iter.util.rararchiveloader import RarArchiveLoaderIterDataPipe as RarArchiveLoader
 from torchdata.datapipes.iter.util.rows2columnar import Rows2ColumnarIterDataPipe as Rows2Columnar
@@ -187,6 +190,7 @@
     "OnlineReader",
     "ParagraphAggregator",
     "ParquetDataFrameLoader",
+    "PinMemory",
     "Prefetcher",
     "RandomSplitter",
     "RarArchiveLoader",

torchdata/datapipes/iter/__init__.pyi.in

@@ -10,6 +10,7 @@ ${init_base}
 from .util.decompressor import CompressionType
 from torchdata._constants import default_timeout_in_s
 from torchdata.datapipes.map import MapDataPipe
+from torchdata.datapipes.utils import pin_memory_fn
 from torch.utils.data import DataChunk, IterableDataset, default_collate
 from torch.utils.data.datapipes._typing import _DataPipeMeta
 from torch.utils.data.datapipes.iter.sharding import SHARDING_PRIORITIES