[aDAG] Support all reduce collective in aDAG

[dengwxn]

Why are these changes needed?

aDAG currently does not support collective APIs. We would like to add support for collective APIs, starting from allreduce.

This PR support allreduce by introducing a syntax sugar ray.experimental.collective.allreduce.bind. The bind accepts arguments input_nodes, op, and transport. It returns a list of CollectiveOutputNode as the allreduce results, with the same size of input_nodes. The allreduce results write to newly allocated tensors. In the COMPUTE operation of CollectiveOutputNode, the corresponding NCCL collective API is called. There are no required changes for the input and output channels of CollectiveOutputNode.

Proposed new API:

import ray.experimental.collective as collective

with InputNode() as inp:
    dag = [worker.return_tensor.bind(inp) for worker in workers]
    dag = collective.allreduce.bind(dag, ReduceOp.SUM)
    dag = MultiOutputNode(dag)

API Requirements:

1. Input nodes are unique.
2. Actor handles are unique.
3. Actor handles match the custom NCCL group if specified.
4. All tensors have the same shape.

Requirements 1-3 are checked in the _CollectiveOperation constructor. Requirement 4 is checked by runtime timeout.

The operation scheduling is also updated to consider the NCCL collective operation. When a NCCL collective node is selected, all the corresponding collective nodes in the collective group should be selected as well.

Related issue number

Meta-issue: #47983

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[stephanie-wang]

Nice work! Left a few remaining comments for cleanups and to see if we can further cut down on the tests that need GPUs to run. The only functionality we really need to test for GPUs is:

• does it provide the expected allreduce result?
• does passing a custom communicator actually result in calling the custom communicator during execution?

Otherwise, all logic should go in unit tests that don't need GPUs to run.