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

python/ray/dag/collective_node.py

@@ -71,6 +71,7 @@ def __init__(
                 raise ValueError(
                     "Expected actor handles to match the custom NCCL group"
                 )
+        self._output_nodes: List[DAGNode] = []
 
     def __str__(self) -> str:
         return (
@@ -88,6 +89,13 @@ def actor_handles(self) -> List["ray.actor.ActorHandle"]:
     def type_hint(self) -> TorchTensorType:
         return self._type_hint
 
+    def _add_output_node(self, output_node: "CollectiveOutputNode"):
+        self._output_nodes.append(output_node)
+
+    @property
+    def output_nodes(self) -> List["CollectiveOutputNode"]:
+        return self._output_nodes
+
     def init_nccl_group(self, nccl_group_id: Optional[str] = None) -> str:
         """
         Initialize the NCCL group if it has not been initialized yet. If `nccl_group_id`
@@ -142,6 +150,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)
@@ -156,14 +172,8 @@ def __init__(
         )
         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_op._add_output_node(self)
+        self.set_requires_nccl_collective(True)
 
     def _copy_impl(
         self,
@@ -188,3 +198,7 @@ def _execute_impl(self, *args, **kwargs):
     @property
     def collective_op(self) -> _CollectiveOperation:
         return self._collective_op
+
+    @property
+    def synchronous_peers(self) -> List["CollectiveOutputNode"]:
+        return self._collective_op.output_nodes