(WIP) [core][compiled graphs] Unify code paths for NCCL P2P and collectives scheduling

python/ray/dag/__init__.py

@@ -17,7 +17,7 @@
     PREV_CLASS_METHOD_CALL_KEY,
     BIND_INDEX_KEY,
     IS_CLASS_METHOD_OUTPUT_KEY,
-    COLLECTIVE_OPERATION_KEY,
+    COLLECTIVE_GROUP_KEY,
     DAGNODE_TYPE_KEY,
 )
 from ray.dag.vis_utils import plot
@@ -38,7 +38,7 @@
     "PREV_CLASS_METHOD_CALL_KEY",
     "BIND_INDEX_KEY",
     "IS_CLASS_METHOD_OUTPUT_KEY",
-    "COLLECTIVE_OPERATION_KEY",
+    "COLLECTIVE_GROUP_KEY",
     "DAGNODE_TYPE_KEY",
     "plot",
     "MultiOutputNode",

python/ray/dag/class_node.py

@@ -180,7 +180,7 @@ def __init__(
             ClassMethodNode
         ] = other_args_to_resolve.get(PREV_CLASS_METHOD_CALL_KEY, None)
         # The index/order when bind() is called on this class method
-        self._bind_index: Optional[int] = other_args_to_resolve.get(
+        self._bind_index: Optional[Union[int, float]] = other_args_to_resolve.get(
             BIND_INDEX_KEY, None
         )
         # Represent if the ClassMethodNode is a class method output. If True,
@@ -252,7 +252,7 @@ def __repr__(self) -> str:
     def get_method_name(self) -> str:
         return self._method_name
 
-    def _get_bind_index(self) -> int:
+    def _get_bind_index(self) -> Union[int, float]:
         return self._bind_index
 
     def _get_remote_method(self, method_name):

python/ray/dag/collective_node.py

@@ -8,17 +8,18 @@
     DAGNode,
     ClassMethodNode,
 )
-from ray.dag.constants import COLLECTIVE_OPERATION_KEY
+from ray.dag.constants import COLLECTIVE_GROUP_KEY
+from ray.dag.sync_group import _SynchronousGroup
 from ray.experimental.channel import ChannelContext
 from ray.experimental.channel.torch_tensor_nccl_channel import _init_nccl_group
 from ray.experimental.channel.torch_tensor_type import GPUCommunicator, TorchTensorType
 from ray.experimental.util.types import _CollectiveOp, ReduceOp
 from ray.util.annotations import DeveloperAPI
 
 
-class _CollectiveOperation:
+class _CollectiveGroup(_SynchronousGroup):
     """
-    Represent metadata for a NCCL collective operation.
+    Represent metadata for a group of actors in a NCCL collective operation.
 
     Args:
         input_nodes: A list of input nodes to the collective operation.
@@ -37,6 +38,8 @@ def __init__(
         op: _CollectiveOp,
         transport: Optional[Union[str, GPUCommunicator]] = None,
     ):
+        super().__init__()
+
         if len(input_nodes) == 0:
             raise ValueError("Expected input nodes for a collective operation")
         if len(set(input_nodes)) != len(input_nodes):
@@ -142,6 +145,14 @@ def __init__(
         method_options: Dict[str, Any],
         other_args_to_resolve: Dict[str, Any],
     ):
+        super().__init__(
+            method_name,
+            method_args,
+            method_kwargs,
+            method_options,
+            other_args_to_resolve,
+        )
+
         # Parse the input node.
         if not (
             isinstance(method_args, tuple)
@@ -151,19 +162,12 @@ def __init__(
             raise ValueError("Expected a single input node")
         self._input_node = method_args[0]
         # Parse the collective operation.
-        self._collective_op: _CollectiveOperation = other_args_to_resolve.get(
-            COLLECTIVE_OPERATION_KEY, None
-        )
-        if self._collective_op is None:
-            raise ValueError("Expected a collective operation")
-
-        super().__init__(
-            method_name,
-            method_args,
-            method_kwargs,
-            method_options,
-            other_args_to_resolve,
+        self._collective_group: _CollectiveGroup = other_args_to_resolve.get(
+            COLLECTIVE_GROUP_KEY, None
         )
+        if self._collective_group is None:
+            raise ValueError("Expected a collective group")
+        self.set_requires_nccl_collective(True)
 
     def _copy_impl(
         self,
@@ -186,5 +190,9 @@ def _execute_impl(self, *args, **kwargs):
         )
 
     @property
-    def collective_op(self) -> _CollectiveOperation:
-        return self._collective_op
+    def collective_group(self) -> _CollectiveGroup:
+        return self._collective_group
+
+    @property
+    def sync_group(self) -> _CollectiveGroup:
+        return self._collective_group