Fix lint for SequentialReadingService and unify the API for process_reset_fn (meta-pytorch#983)

ejguan · facebook-github-bot · commit 01fc76200354 · 2023-02-03T08:53:56.000-08:00
Summary: This is the follow up for 807db8f ### Changes - Fix lint - Fix checkpoint result - Add inline doc - Unify the API for process_reset_fn and dispatch_process_reset_fn After this PR, I will land a test case and tutorial for SequentialRS Pull Request resolved: meta-pytorch#983 Reviewed By: NivekT Differential Revision: D42979577 Pulled By: ejguan fbshipit-source-id: 1cb601c20701cf44a1f7b3afe972a9436cdb9465
diff --git a/docs/source/dataloader2.rst b/docs/source/dataloader2.rst
@@ -33,6 +33,7 @@ ReadingService
     DistributedReadingService
     MultiProcessingReadingService
     PrototypeMultiProcessingReadingService
+    SequentialReadingService
 
 Each ``ReadingServices`` would take the ``DataPipe`` graph and rewrite it to achieve a few features like dynamic sharding, sharing random seeds and snapshoting for multi-/distributed processes. For more detail about those features, please refer to `the documentation <reading_service.html>`_.
 
diff --git a/torchdata/dataloader2/__init__.py b/torchdata/dataloader2/__init__.py
@@ -13,7 +13,7 @@
     MultiProcessingReadingService,
     PrototypeMultiProcessingReadingService,
     ReadingServiceInterface,
-    SequentialReadingService
+    SequentialReadingService,
 )
 from torchdata.dataloader2.shuffle_spec import ShuffleSpec
 
diff --git a/torchdata/dataloader2/reading_service.py b/torchdata/dataloader2/reading_service.py
@@ -64,7 +64,9 @@ def finalize(self) -> None:
         """
         pass
 
-    def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None) -> Optional[Callable[[DataPipe], DataPipe]]:
+    def initialize_iteration(
+        self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None
+    ) -> Optional[Callable[[DataPipe], DataPipe]]:
         r"""
         ``ReadingService`` spins up service for an epoch. Called at the beginning
         of every time getting ``DataLoader2`` iterator.
@@ -73,6 +75,11 @@ def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Opt
             seed_generator: SeedGenerator object created and managed by DataLoader2. As the single
                 source of randomness, it will governs the determinism for all of random operations
                 with the graph of DataPipes.
+            iter_reset_fn: Optional reset function from the prior ``ReadingServcie``
+                when ``SequentialReadingService`` chains multiple ``ReadingServices``
+
+        Returns:
+            A new ``iter_reset_fn`` to be used by subseqeuent ``ReadingService``
 
         Example:
             MultiProcessingReadingService starts setting worker seeds per process and prefetching
@@ -303,7 +310,9 @@ def initialize(self, datapipe: DataPipe) -> DataPipe:
 
         return self._end_datapipe  # type: ignore[return-value]
 
-    def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None) -> Optional[Callable[[DataPipe], DataPipe]]:
+    def initialize_iteration(
+        self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None
+    ) -> Optional[Callable[[DataPipe], DataPipe]]:
         if self._pg is not None:
             shared_seed_int = dist_share_seed(seed_generator.generate_shared_seed(), self._pg)
             seed_generator.seed(shared_seed_int)
@@ -318,7 +327,9 @@ def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Opt
                 # Stop prefetching first
                 self._main_prefetch_datapipe.reset()  # type: ignore[union-attr]
             # Send the shared seed to subprocesses
-            call_on_epoch_reset = partial(process_reset_fn, custom_reset_fn=self.worker_reset_fn, custom_dispatch_process_reset_fn=iter_reset_fn)
+            call_on_epoch_reset = partial(
+                process_reset_fn, iter_reset_fn=iter_reset_fn, custom_reset_fn=self.worker_reset_fn
+            )
             assert self._worker_consumer_datapipe is not None
             self._worker_consumer_datapipe.reset_epoch(call_on_epoch_reset, seed_generator)
         # In-process (num_workers == 0)
@@ -328,6 +339,7 @@ def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Opt
             # (random, torch and numpy), if users have already seeded them in the main process
             # TODO(ejguan): This should be fixed by adding a method to isolate global RNGs
             pass
+        return None
 
     def __del__(self):
         self.finalize()
@@ -485,7 +497,9 @@ def initialize(self, datapipe: DataPipe) -> DataPipe:
         self._datapipe = datapipe
         return datapipe
 
-    def initialize_iteration(self, seed_generator: SeedGenerator) -> None:
+    def initialize_iteration(
+        self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None
+    ) -> Optional[Callable[[DataPipe], DataPipe]]:
         r"""
         Shares the same seed from rank 0 to other ranks across the distributed processes
         and apply the random seed to the ``DataPipe`` graph.
@@ -496,6 +510,7 @@ def initialize_iteration(self, seed_generator: SeedGenerator) -> None:
         seed_generator.seed(shared_seed)
         seed_generator = seed_generator.spawn(self._rank, inplace=True)
         set_graph_random_seed(self._datapipe, seed_generator)
+        return None
 
     def __del__(self):
         self.finalize()
@@ -509,7 +524,7 @@ def finalize(self) -> None:
             self._pg = None
 
 
-class SequentialReadingService(ReadingServiceInterface):
+class SequentialReadingService(CheckpointableReadingServiceInterface):
     def __init__(self, *reading_services):
         self.reading_services = reading_services
 
@@ -525,10 +540,15 @@ def finalize(self) -> None:
             rs.finalize()
 
     # Sequential Order
-    def initialize_iteration(self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None) -> Optional[Callable[[DataPipe], DataPipe]]:
+    def initialize_iteration(
+        self, seed_generator: SeedGenerator, iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None
+    ) -> Optional[Callable[[DataPipe], DataPipe]]:
         chained_iter_reset_fn = iter_reset_fn
         for rs in self.reading_services:
-            chained_iter_reset_fn = rs.initialize_iteration(seed_generator=seed_generator, iter_reset_fn=chained_iter_reset_fn)
+            chained_iter_reset_fn = rs.initialize_iteration(
+                seed_generator=seed_generator, iter_reset_fn=chained_iter_reset_fn
+            )
+        return None
 
     # Reversed Order
     def finalize_iteration(self) -> None:
@@ -540,9 +560,12 @@ def checkpoint(self) -> bytes:
         states = []
         for rs in self.reading_services:
             states.append(rs.checkpoint())
+        return b"\n".join(states)
 
     # Sequential Order, to align with initialize
-    def restore(self, datapipe, serialized_state) -> DataPipe:
-        for rs, state in zip(self.reading_service, serialized_state):
+    def restore(self, datapipe, serialized_state: bytes) -> DataPipe:
+        states = serialized_state.split(b"\n")
+        assert len(states) == len(self.reading_services)
+        for rs, state in zip(self.reading_services, states):
             datapipe = rs.restore(datapipe, state)
         return datapipe
diff --git a/torchdata/dataloader2/utils/worker.py b/torchdata/dataloader2/utils/worker.py
@@ -116,7 +116,8 @@ def dispatch_process_reset_fn(
     datapipe: DataPipe,
     worker_info: WorkerInfo,
     seed_generator: SeedGenerator,
-    custom_dispatch_process_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None,
+    iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None,
+    custom_reset_fn: Optional[Callable[[DataPipe, WorkerInfo, SeedGenerator], DataPipe]] = None,
 ) -> DataPipe:
     r"""
     Based on the distributed shared random seed, this function is used to set the random state
@@ -131,8 +132,12 @@ def dispatch_process_reset_fn(
     dps = list_dps(graph)
     set_datapipes_seed(dps, seed_generator=seed_generator, distributed_shared=True)
 
-    if custom_dispatch_process_reset_fn is not None:
-        datapipe = custom_dispatch_process_reset_fn(datapipe)
+    if iter_reset_fn is not None:
+        datapipe = iter_reset_fn(datapipe)
+        assert isinstance(datapipe, (IterDataPipe, MapDataPipe))
+
+    if custom_reset_fn is not None:
+        datapipe = custom_reset_fn(datapipe, worker_info, seed_generator)
         assert isinstance(datapipe, (IterDataPipe, MapDataPipe))
 
     return datapipe
@@ -142,8 +147,8 @@ def process_reset_fn(
     datapipe: DataPipe,
     worker_info: WorkerInfo,
     seed_generator: SeedGenerator,
-    custom_reset_fn: Optional[Callable[[DataPipe, WorkerInfo], DataPipe]] = None,
-    custom_dispatch_process_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None,
+    iter_reset_fn: Optional[Callable[[DataPipe], DataPipe]] = None,
+    custom_reset_fn: Optional[Callable[[DataPipe, WorkerInfo, SeedGenerator], DataPipe]] = None,
 ) -> DataPipe:
     r"""
     Based on the distributed shared random seed and worker id, this function is used to
@@ -160,14 +165,20 @@ def process_reset_fn(
         # Only send the reset epoch message once
         if worker_info.worker_id == 0:
             # Use WorkerInfo(1, 0)
-            dispatch_reset_fn = partial(dispatch_process_reset_fn, custom_dispatch_process_reset_fn=custom_dispatch_process_reset_fn)
+            dispatch_reset_fn = partial(
+                dispatch_process_reset_fn, iter_reset_fn=iter_reset_fn, custom_reset_fn=custom_reset_fn
+            )
             dispatch_process_consumer_dp.reset_epoch(dispatch_reset_fn, seed_generator)
 
     # Set global random states
     _set_global_random_state(seed_generator)
 
     set_graph_random_seed(datapipe, seed_generator)
 
+    if iter_reset_fn is not None:
+        datapipe = iter_reset_fn(datapipe)
+        assert isinstance(datapipe, (IterDataPipe, MapDataPipe))
+
     if custom_reset_fn is not None:
         datapipe = custom_reset_fn(datapipe, worker_info, seed_generator)
         assert isinstance(datapipe, (IterDataPipe, MapDataPipe))

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@`
`13`	`13`	`MultiProcessingReadingService,`
`14`	`14`	`PrototypeMultiProcessingReadingService,`
`15`	`15`	`ReadingServiceInterface,`
`16`		`- SequentialReadingService`
	`16`	`+ SequentialReadingService,`
`17`	`17`	`)`
`18`	`18`	`from torchdata.dataloader2.shuffle_spec import ShuffleSpec`
`19`	`19`