[Data Loading] Round-based per-epoch shuffling data loader for distributed training.

[clarkzinzow]

A round-based per-epoch shuffling data loader for distributed training. I'm opening this PR early in order to facilitate easier collaboration.

[richardliaw]

I got it to work on 4 GPUs.

Nit. Seeing this warning:

2021-04-28 08:04:47,389 WARNING import_thread.py:133 -- The remote function 'ray.experimental.data_loader.shuffle.consume' has been exported 100 times. It's possible that this warning is accidental, but this may indicate that the same remote function is being defined repeatedly from within many tasks and exported to all of the workers. This can be a performance issue and can be resolved by defining the remote function on the driver instead. See #6240 for more discussion.

[richardliaw]

Example working here: https://gist.github.com/9040abcc654ce6b5ed817b8263d723e2

I also had to add this to get the workload to work:

diff --git a/python/ray/experimental/data_loader/multiqueue.py b/python/ray/experimental/data_loader/multiqueue.py
index 9e0cb0394..4b70f1c2e 100644
--- a/python/ray/experimental/data_loader/multiqueue.py
+++ b/python/ray/experimental/data_loader/multiqueue.py
@@ -3,7 +3,9 @@ from typing import Optional, Any, List, Dict
 from collections.abc import Iterable

 import ray
-
+import logging
+import time
+logger = logging.getLogger(__name__)

 class Empty(Exception):
     pass
@@ -49,13 +51,19 @@ class MultiQueue:

     def __init__(self, num_queues: int, maxsize: int = 0,
                  name: str = None, connect: bool = False,
-                 actor_options: Optional[Dict] = None) -> None:
+                 actor_options: Optional[Dict] = None, retries=5) -> None:
         self.num_queues = num_queues
         self.maxsize = maxsize
         if connect:
             assert actor_options is None
             assert name is not None
-            self.actor = ray.get_actor(name)
+            for i in range(retries):
+                try:
+                    self.actor = ray.get_actor(name)
+                except ValueError:
+                    logger.info(
+                        f"Did not acquire actor. Trying again [{i}/{retries}.")
+                    time.sleep(3)
         else:
             actor_options = actor_options or {}
             if name is not None:

[richardliaw]

Trying now on 16 GPU cluster with the provided example. A couple notes:

• Sometimes I see a hang. Not sure why.
• Not sure how to set num_reducers and num_mappers. It'd be great to have some good defaults so that users won't need to touch that (usually this will be scaffolded away so deeply that users can't touch it anyways).
• You can use this cluster: https://gist.github.com/c8bf8fd7eab92d89ea71c612377ad29d

ray up -y cluster
ray submit cluster ray_torch_shuffle.py

[stephanie-wang]

Just reviewed shuffle.py mainly.

[clarkzinzow]

Moved to an external repo.