[Collective] Ray CPU collectives now available (ray-project#14277)

Co-authored-by: YLJALDC <[email protected]> Co-authored-by: Ezra-H <[email protected]> Co-authored-by: Ezra-H <[email protected]>
sihanwang41 · Mar 9, 2021 · 2505bc8 · 2505bc8
1 parent ed89354
commit 2505bc8
Show file tree

Hide file tree

Showing 36 changed files with 1,797 additions and 81 deletions.
diff --git a/python/ray/util/collective/collective.py b/python/ray/util/collective/collective.py
@@ -7,15 +7,28 @@
 import ray
 from ray.util.collective import types
 
-_GLOO_AVAILABLE = False
 _NCCL_AVAILABLE = True
+_GLOO_AVAILABLE = True
+
+logger = logging.getLogger(__name__)
 
 try:
- from ray.util.collective.collective_group import NCCLGroup
+ from ray.util.collective.collective_group.\
+ nccl_collective_group import NCCLGroup
 except ImportError:
  _NCCL_AVAILABLE = False
+ logger.warning("NCCL seems unavailable. Please install Cupy "
+ "following the guide at: "
+ "https://docs.cupy.dev/en/stable/install.html.")
 
-logger = logging.getLogger(__name__)
+try:
+ from ray.util.collective.collective_group.\
+ gloo_collective_group import GLOOGroup
+except ImportError:
+ _GLOO_AVAILABLE = False
+ logger.warning("PyGloo seems unavailable. Please install PyGloo "
+ "following the guide at: "
+ "https://github.com/ray-project/pygloo.")
 
 
 def nccl_available():
@@ -48,7 +61,15 @@ def create_collective_group(self, backend, world_size, rank, group_name):
  if backend == types.Backend.MPI:
  raise RuntimeError("Ray does not support MPI.")
  elif backend == types.Backend.GLOO:
- raise NotImplementedError()
+ logger.debug("Creating GLOO group: '{}'...".format(group_name))
+ g = GLOOGroup(
+ world_size,
+ rank,
+ group_name,
+ store_type="redis",
+ device_type="tcp")
+ self._name_group_map[group_name] = g
+ self._group_name_map[g] = group_name
  elif backend == types.Backend.NCCL:
  logger.debug("Creating NCCL group: '{}'...".format(group_name))
  g = NCCLGroup(world_size, rank, group_name)

diff --git a/python/ray/util/collective/collective_group/__init__.py b/python/ray/util/collective/collective_group/__init__.py
@@ -1,3 +0,0 @@
-from .nccl_collective_group import NCCLGroup
-
-__all__ = ["NCCLGroup"]