Enable miniepoch for MultiProcessingReadingService (pytorch#1170)

Summary:
A retry of pytorch#1147
- Add pause, resume, limit API to MultiProcessingReadingService to support miniepoch API.
- When limit is requested, `IterQueueWrapper` will make sure only `limit` number of requests are sent to worker process and dispatching proccess.
- Consolidate a few API/messages
  - reset and reset_epoch is merged as a single message
  - Process reset function is shared for both worker process and dispatching process
- For pause/resume/limit/reset_epoch, add a counter to sync between loops to make sure the DataPipe graph is in-place modified once.
- Remove unused thread eventloop

Pull Request resolved: pytorch#1170

Differential Revision: D45666912

Pulled By: ejguan

fbshipit-source-id: 1e88dde3af0f406c062ca523e7b96a3c79a74b6d

Author: ejguan

test/dataloader2/test_dataloader2.py

@@ -23,7 +23,6 @@
 from torch.utils.data.datapipes.iter.sharding import SHARDING_PRIORITIES
 
 from torchdata.dataloader2 import (
-    communication,
     DataLoader2,
     DistributedReadingService,
     MultiProcessingReadingService,
@@ -35,7 +34,6 @@
 from torchdata.dataloader2.graph import DataPipe, list_dps, replace_dp, set_datapipes_seed, traverse_dps
 from torchdata.dataloader2.random import SeedGenerator
 from torchdata.datapipes.iter import IterableWrapper, IterDataPipe, ShardingRoundRobinDispatcher
-from torchdata.datapipes.map import SequenceWrapper
 
 try:
     import dill
@@ -259,73 +257,6 @@ def test_dataloader2_shuffle(self) -> None:
         pass
 
 
-@unittest.skipIf(
-    TEST_WITH_TSAN,
-    "Fails with TSAN with the following error: starting new threads after multi-threaded "
-    "fork is not supported. Dying (set die_after_fork=0 to override)",
-)
-class TestDataLoader2EventLoop(TestCase):
-    # TODO: This needs fixing, see issue 624
-    # @skipIfNoDill
-    # def test_basic_threading(self):
-    #     def clean_me(process, req_queue, res_queue):
-    #         req_queue.put(communication.messages.TerminateRequest())
-    #         _ = res_queue.get()
-    #         process.join()
-    #
-    #     it = list(range(100))
-    #     numbers_dp = IterableWrapper(it)
-    #     (process, req_queue, res_queue, _thread_local_datapipe) = communication.eventloop.CreateThreadForDataPipeline(numbers_dp)
-    #
-    #     process.start()
-    #     local_datapipe = communication.iter.QueueWrapper(
-    #         communication.protocol.IterDataPipeQueueProtocolClient(req_queue, res_queue))
-    #
-    #     actual = list(local_datapipe)
-    #     clean_me(process, req_queue, res_queue)
-    #
-    #     self.assertEqual(list(range(100)), actual)
-
-    @skipIfNoDill
-    def test_basic_mapdatapipe_threading(self):
-        def clean_me(process, req_queue, res_queue):
-            req_queue.put(communication.messages.TerminateRequest())
-            _ = res_queue.get()
-            process.join()
-
-        input_len = 100
-        it = list(range(input_len))
-        numbers_dp = SequenceWrapper(it)
-        (process, req_queue, res_queue, _thread_local_datapipe) = communication.eventloop.CreateThreadForDataPipeline(
-            numbers_dp,
-            thread_name="worker thread",
-        )
-
-        process.start()
-
-        # Functional Test: Ensure that you can retrieve every element from the Queue and DataPipe
-        local_datapipe = communication.map.QueueWrapperForMap(
-            communication.protocol.MapDataPipeQueueProtocolClient(req_queue, res_queue)
-        )
-        actual = list(local_datapipe)
-        self.assertEqual([(x, x) for x in range(100)], actual)
-
-        # Functional Test: raise Error when input
-        local_datapipe = communication.map.QueueWrapperForMap(
-            communication.protocol.MapDataPipeQueueProtocolClient(req_queue, res_queue)
-        )
-        with self.assertRaisesRegex(IndexError, "out of bound"):
-            local_datapipe[1000]
-
-        # __len__ Test: Ensure that the correct length is returned
-        local_datapipe = communication.map.QueueWrapperForMap(
-            communication.protocol.MapDataPipeQueueProtocolClient(req_queue, res_queue)
-        )
-        self.assertEqual(input_len, len(local_datapipe))
-
-        clean_me(process, req_queue, res_queue)
-
-
 def _x_mult_2(d):
     return d * 2
 

test/dataloader2/test_mprs.py

@@ -40,6 +40,7 @@ def _non_dispatching_dp(n_elements=1000):
 
 def _dispatching_dp(n_elements=1000):
     dp = IterableWrapper(list(range(n_elements))).shuffle()
+    dp = dp.prefetch(20)
     dp = dp.sharding_round_robin_dispatch(SHARDING_PRIORITIES.MULTIPROCESSING)
     dp = dp.map(_add_one).batch(16)
     return dp

torchdata/dataloader2/communication/eventloop.py

@@ -4,11 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import pickle
-import threading
 import time
 
 from itertools import zip_longest
+from typing import Dict, List
 
 import torch
 
@@ -32,36 +31,42 @@
     "DataPipeToQueuesLoop",
     "CreateProcessForDataPipeline",
     "CreateProcessForMultipleDataPipelines",
-    "CreateThreadForDataPipeline",
 ]
 
 
-class _ResetCounter:
+class _RequestCounter:
     exp_cnt: int
-    cnt: int
-    _reached: bool
+    _keys: List[str] = ["limit", "pause", "reset_epoch", "resume"]
+    _cnt: Dict[str, int]
+    _reached: Dict[str, bool]
 
     def __init__(self, exp_cnt: int):
         self.exp_cnt = exp_cnt
-        self.cnt = 0
-        self._reached = False
-
-    def increment(self) -> None:
-        self.cnt += 1
-        assert self.cnt <= self.exp_cnt
-
-    def is_reached(self) -> bool:
-        if self.cnt == self.exp_cnt:
-            self._reached = True
-        return self._reached
-
-    def reset(self) -> None:
-        if self._reached:
-            self._reached = False
-            self.cnt = 0
-
-
-def MultipleDataPipesToQueuesLoop(source_datapipes, req_queues, res_queues, process_name, call_on_process_init=None):
+        self._cnt = {k: 0 for k in self._keys}
+        self._reached = {k: False for k in self._keys}
+
+    def increment(self, key: str) -> None:
+        assert key in self._reached
+        self._cnt[key] += 1
+        assert self._cnt[key] <= self.exp_cnt
+        if self._cnt[key] == self.exp_cnt:
+            self._reached[key] = True
+
+    def is_reached(self, key: str) -> bool:
+        assert key in self._reached
+        return self._reached[key]
+
+    def reset(self, key: str) -> None:
+        assert key in self._reached and self._reached[key]
+        assert self._cnt[key] >= 1
+        self._cnt[key] -= 1
+        if self._cnt[key] == 0:
+            self._reached[key] = False
+
+
+def MultipleDataPipesToQueuesLoop(
+    source_datapipes, req_queues, res_queues, process_name, worker_info, call_on_process_init=None, custom_reset_fn=None
+):
     r"""
     Set the appropriate pipes and protocol server type, and create a loop over multiple datapipes
     with the protocol server in a non-blocking manner.
@@ -71,7 +76,9 @@ def MultipleDataPipesToQueuesLoop(source_datapipes, req_queues, res_queues, proc
         req_queue: Multiprocessing queue providing requests from the worker process
         res_queue: Multiprocessing queue sending results to the worker process
         process_name: The name of process (used for logging and exception handling)
+        worker_info: Worker information (worker id and number of workers)
         call_on_process_init: Not allowed by dispatching process for now.
+        custom_reset_fn: Optional callable function to reset the DataPipe.
     """
     assert call_on_process_init is None, "``MultipleDataPipesToQueuesLoop`` does not support call_on_process_init"
     num_loops = len(source_datapipes)
@@ -82,21 +89,24 @@ def MultipleDataPipesToQueuesLoop(source_datapipes, req_queues, res_queues, proc
     torch.set_num_threads(1)
 
     loops = []
-    reset_iterator_counter = _ResetCounter(num_loops)
+    request_counter = _RequestCounter(num_loops)
 
+    loop_id = 0
     for source_datapipe, req_queue, res_queue in zip(source_datapipes, req_queues, res_queues):
-        # Extract Serialization Wrapper
-        source_datapipe = extract_wrapper(source_datapipe)
         loops.append(
             _create_datapipe_queue_loop(
                 source_datapipe,
                 req_queue,
                 res_queue,
                 process_name,
+                loop_id,
+                worker_info,
+                custom_reset_fn,
                 blocking_request_get=False,
-                reset_iterator_counter=reset_iterator_counter,
+                request_counter=request_counter,
             )
         )  # Non-blocking request with reset counters
+        loop_id += 1
 
     # Using `zip_longest` to guarantee the process is terminated only when
     # all loops have received `TerminateRequest`
@@ -107,7 +117,9 @@ def MultipleDataPipesToQueuesLoop(source_datapipes, req_queues, res_queues, proc
         time.sleep(0)
 
 
-def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue, process_name, call_on_process_init=None):
+def DataPipeToQueuesLoop(
+    source_datapipe, req_queue, res_queue, process_name, worker_info, call_on_process_init=None, custom_reset_fn=None
+):
     r"""
     Initialize with the given init function, set the appropriate pipe and protocol server type, and
     create a loop with the protocol server.
@@ -117,8 +129,10 @@ def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue, process_name, ca
         req_queue: Multiprocessing queue providing requests from the main process
         res_queue: Multiprocessing queue sending results to the main process
         process_name: The name of process (used for logging and exception handling)
+        worker_info: Worker information (worker id and number of workers)
         call_on_process_init: Callable function will be called at the time of worker process initialization.
             Users can provide it to modify the DataPipe grpah in the worker process.
+        custom_reset_fn: Optional callable function to reset the DataPipe.
     """
     # Extract Serialization Wrapper
     source_datapipe = extract_wrapper(source_datapipe)
@@ -128,7 +142,16 @@ def DataPipeToQueuesLoop(source_datapipe, req_queue, res_queue, process_name, ca
 
     torch.set_num_threads(1)
 
-    loop = _create_datapipe_queue_loop(source_datapipe, req_queue, res_queue, process_name, blocking_request_get=True)
+    loop = _create_datapipe_queue_loop(
+        source_datapipe,
+        req_queue,
+        res_queue,
+        process_name,
+        worker_info.worker_id,
+        worker_info,
+        custom_reset_fn,
+        blocking_request_get=True,
+    )
 
     for _ in loop:
         pass
@@ -139,8 +162,11 @@ def _create_datapipe_queue_loop(
     req_queue,
     res_queue,
     process_name,
+    loop_id,
+    worker_info,
+    custom_reset_fn=None,
     blocking_request_get=True,
-    reset_iterator_counter=None,
+    request_counter=None,
 ):
     if isinstance(source_datapipe, IterDataPipe):
         pipe_type = communication.iter
@@ -155,51 +181,33 @@ def _create_datapipe_queue_loop(
         source_datapipe,
         protocol_type(req_queue, res_queue),
         process_name=process_name,
+        loop_id=loop_id,
+        worker_info=worker_info,
+        custom_reset_fn=custom_reset_fn,
         blocking_request_get=blocking_request_get,
-        reset_iterator_counter=reset_iterator_counter,
+        request_counter=request_counter,
     )
 
 
-def CreateProcessForDataPipeline(multiprocessing_ctx, datapipe, process_name, call_on_process_init=None):
+def CreateProcessForDataPipeline(
+    multiprocessing_ctx, datapipe, process_name, worker_info, call_on_process_init=None, custom_reset_fn=None
+):
     r"""
     Given a DataPipe, creates a new process with ``DataPipeToQueuesLoop`` as target,
     and returns ``(process, req_queue, res_queue)``.
     """
     req_queue = multiprocessing_ctx.Queue()
     res_queue = multiprocessing_ctx.Queue()
     process = multiprocessing_ctx.Process(
-        target=DataPipeToQueuesLoop, args=(datapipe, req_queue, res_queue, process_name, call_on_process_init)
+        target=DataPipeToQueuesLoop,
+        args=(datapipe, req_queue, res_queue, process_name, worker_info, call_on_process_init, custom_reset_fn),
     )
     return process, req_queue, res_queue
 
 
-def CreateThreadForDataPipeline(datapipe, thread_name):
-    r"""
-    Given a DataPipe, creates a copy of the DataPipe, starts a new Thread with ``DataPipeToQueuesLoop`` as target,
-    and returns ``(process, req_queue, res_queue, new_copied_datapipe)``.
-    """
-    req_queue = communication.queue.ThreadingQueue()
-    res_queue = communication.queue.ThreadingQueue()
-
-    try:
-        new_datapipe = pickle.loads(pickle.dumps(datapipe))
-    except Exception as pe:
-        if HAS_DILL:
-            try:
-                new_datapipe = dill.loads(dill.dumps(datapipe))
-            except Exception as de:
-                raise Exception("Unable to dill DataPipe to make thread local copy", de)
-
-        else:
-            raise Exception("Unable to pickle DataPipe to make thread local copy (consider installing `dill`)", pe)
-
-    process = threading.Thread(
-        target=DataPipeToQueuesLoop, args=(new_datapipe, req_queue, res_queue, thread_name), daemon=True
-    )
-    return process, req_queue, res_queue, new_datapipe
-
-
-def CreateProcessForMultipleDataPipelines(multiprocessing_ctx, datapipes, process_name):
+def CreateProcessForMultipleDataPipelines(
+    multiprocessing_ctx, datapipes, process_name, worker_info, custom_reset_fn=None
+):
     r"""
     Given a DataPipe, creates a new process with ``MultipleDataPipesToQueuesLoop`` as target,
     and returns ``(process, [req_queue_0, ...], [res_queue_0, ...])``.
@@ -211,6 +219,7 @@ def CreateProcessForMultipleDataPipelines(multiprocessing_ctx, datapipes, proces
         res_queues.append(multiprocessing_ctx.Queue())
 
     process = multiprocessing_ctx.Process(
-        target=MultipleDataPipesToQueuesLoop, args=(datapipes, req_queues, res_queues, process_name)
+        target=MultipleDataPipesToQueuesLoop,
+        args=(datapipes, req_queues, res_queues, process_name, worker_info, custom_reset_fn),
     )
     return process, req_queues, res_queues