Repository: ray-project/ray
Issue category: Bug Report

Issue:
[core][compiled graphs] Run fails only when using NCCL channels
### What happened + What you expected to happen

The following compiled graph works when I disable NCCL channel transfer: https://github.com/anmscale/verl/blob/ray_cg/verl/trainer/ppo/ray_trainer.py#L621-L662

However, when I enable NCCL transfer in these two steps I get an error:
https://github.com/anmscale/verl/blob/ray_cg/verl/trainer/ppo/ray_trainer.py#L638-L652

There are two different errors:
1) When I type-hint output of Actor-0 to Actor-0, which shouldn't be allowed at compile time but somehow works. I get the following error:
```
(main_task pid=2531538) Waiting for worker tasks to exit
(main_task pid=2531538) Teardown complete
Traceback (most recent call last):
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 97, in main
    ray.get(main_task.remote(config))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 2763, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 914, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::main_task() (pid=2531538, ip=10.0.0.27)
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 195, in main_task
    trainer.fit()
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/ray_trainer.py", line 733, in fit
    metrics = self.run_compiled_graph(batch_dict)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/ray_trainer.py", line 699, in run_compiled_graph
    dag_metrics = ray.get(dag_futures, timeout=100000)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 106, in get
    return _process_return_vals(return_vals, True)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 22, in _process_return_vals
    raise val.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::WorkerDict.__ray_call__() (pid=2531917, ip=10.0.0.27)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/single_controller/ray/base.py", line 435, in func
    return getattr(self.worker_dict[key], name)(*args, **kwargs)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/single_controller/base/decorator.py", line 406, in inner
    return func(*args, **kwargs)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/trainer/ppo/workers/fsdp_workers.py", line 926, in compute_scores_and_advantage
    data = DataProto.concat([data1, data2, data3, data4])
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_20d7fa7e64594a715b9685aa68d7af73884e34d9/verl/protocol.py", line 430, in concat
    batch_lst.append(batch.batch)
AttributeError: 'ray._raylet.MessagePackSerializedObject' object has no attribute 'batch'

```
2- When I enable nccl only across different Actors (i.e., disable it within Actor-0) I get the following:
```
(main_task pid=2794522) Tearing down compiled DAG
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 21547c8b679f8af158377dd190000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 6106071252e14600d3031a3290000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, 0c16245c6e3308c33ac657f390000000)
(main_task pid=2794522) Cancelling compiled worker on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, ffe689ca0de0981de9b764eb90000000)
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(main_task pid=2794522) Exception ignored in: <function CompiledDAGRef.__del__ at 0x71bb6f199c10>
(main_task pid=2794522) Traceback (most recent call last):
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 93, in __del__
(main_task pid=2794522)     self.get()
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
(main_task pid=2794522)     return_vals = self._dag._execute_until(
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2122, in _execute_until
(main_task pid=2794522)     return self._get_execution_results(execution_index, channel_index)
(main_task pid=2794522)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2047, in _get_execution_results
(main_task pid=2794522)     assert execution_index in self._result_buffer
(main_task pid=2794522) AssertionError: 
(WorkerDict pid=2795038) /home/ray/anaconda3/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py:689: FutureWarning: FSDP.state_dict_type() and FSDP.set_state_dict_type() are being deprecated. Please use APIs, get_state_dict() and set_state_dict(), which can support different parallelisms, FSDP1, FSDP2, DDP. API doc: https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.state_dict.get_state_dict .Tutorial: https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html . [repeated 3x across cluster]
(WorkerDict pid=2795038)   warnings.warn( [repeated 3x across cluster]
(WorkerDict pid=2794846) TIMER: compute_values, CUDA time (s): 5.9334, perf_counter time (s): 5.9336
(WorkerDict pid=2794846) TIMER: union, CUDA time (s): 0.0001, perf_counter time (s): 0.0002
(WorkerDict pid=2794846) ERROR:2024-12-04 09:11:45,151:Compiled DAG task exited with exception
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 161, in do_exec_tasks
(WorkerDict pid=2794846)     done = tasks[operation.exec_task_idx].exec_operation(
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 644, in exec_operation
(WorkerDict pid=2794846)     return self._read(overlap_gpu_communication)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 537, in _read
(WorkerDict pid=2794846)     input_data = self.input_reader.read()
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 296, in read
(WorkerDict pid=2794846)     outputs = self._read_list(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 321, in _read_list
(WorkerDict pid=2794846)     results.append(c.read(timeout))
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/cached_channel.py", line 88, in read
(WorkerDict pid=2794846)     value = self._inner_channel.read(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 751, in read
(WorkerDict pid=2794846)     return self._channel_dict[actor_id].read(timeout)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/intra_process_channel.py", line 64, in read
(WorkerDict pid=2794846)     return ctx.get_data(self._channel_id)
(WorkerDict pid=2794846)   File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/serialization_context.py", line 47, in get_data
(WorkerDict pid=2794846)     assert (
(WorkerDict pid=2794846) AssertionError: Channel 9d1c4b93-dd26-4f96-803d-2231c70dff4e does not exist in the buffer.
(WorkerDict pid=2795037) Destructing NCCL group on actor: Actor(create_colocated_worker_cls.<locals>.WorkerDict, ffe689ca0de0981de9b764eb90000000)
(main_task pid=2794522) Waiting for worker tasks to exit
(main_task pid=2794522) Teardown complete
Error executing job with overrides: ['data.train_files=/mnt/cluster_storage/data/gsm8k/train.parquet', 'data.val_files=/mnt/cluster_storage/data/gsm8k/test.parquet', 'data.train_batch_size=256', 'data.val_batch_size=512', 'data.max_prompt_length=512', 'data.max_response_length=512', 'actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat', 'actor_rollout_ref.actor.optim.lr=1e-6', 'actor_rollout_ref.actor.ppo_mini_batch_size=64', 'actor_rollout_ref.actor.ppo_micro_batch_size=8', 'actor_rollout_ref.actor.fsdp_config.param_offload=False', 'actor_rollout_ref.actor.fsdp_config.grad_offload=False', 'actor_rollout_ref.actor.fsdp_config.optimizer_offload=False', 'actor_rollout_ref.rollout.log_prob_micro_batch_size=32', 'actor_rollout_ref.rollout.tensor_model_parallel_size=4', 'actor_rollout_ref.rollout.name=vllm', 'actor_rollout_ref.rollout.gpu_memory_utilization=0.4', 'actor_rollout_ref.ref.log_prob_micro_batch_size=32', 'actor_rollout_ref.ref.fsdp_config.param_offload=True', 'critic.optim.lr=1e-5', 'critic.model.path=deepseek-ai/deepseek-llm-7b-chat', 'critic.model.enable_gradient_checkpointing=False', 'critic.ppo_micro_batch_size=8', 'critic.model.fsdp_config.param_offload=False', 'critic.model.fsdp_config.grad_offload=False', 'critic.model.fsdp_config.optimizer_offload=False', 'algorithm.kl_ctrl.kl_coef=0.001', 'trainer.critic_warmup=0', 'trainer.logger=[console]', 'trainer.project_name=verl_example_gsm8k', 'trainer.experiment_name=deepseek_llm_7b_function_rm', 'trainer.n_gpus_per_node=4', 'trainer.nnodes=1', 'trainer.save_freq=-1', 'trainer.test_freq=10', 'trainer.total_epochs=5', 'trainer.default_hdfs_dir=/mnt/local_storage/experiments/gsm8k/ppo/deepseek_llm_7b_function_rm', 'trainer.use_rcg=True']
Traceback (most recent call last):
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 97, in main
    ray.get(main_task.remote(config))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 2763, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/_private/worker.py", line 914, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(RayChannelError): ray::main_task() (pid=2794522, ip=10.0.0.27)
  File "/home/ray/default/verl/verl/trainer/main_ppo.py", line 195, in main_task
    trainer.fit()
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_326f52a2699eb2dc53e27c8a8ba5a9863db795ec/verl/trainer/ppo/ray_trainer.py", line 733, in fit
    metrics = self.run_compiled_graph(batch_dict)
  File "/tmp/ray/session_2024-11-27_18-08-58_172807_6952/runtime_resources/working_dir_files/_ray_pkg_326f52a2699eb2dc53e27c8a8ba5a9863db795ec/verl/trainer/ppo/ray_trainer.py", line 699, in run_compiled_graph
    dag_metrics = ray.get(dag_futures, timeout=100000)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/compiled_dag_ref.py", line 103, in get
    return_vals = self._dag._execute_until(
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/dag/compiled_dag_node.py", line 2115, in _execute_until
    self._dag_output_fetcher.read(timeout),
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 296, in read
    outputs = self._read_list(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/common.py", line 321, in _read_list
    results.append(c.read(timeout))
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 751, in read
    return self._channel_dict[actor_id].read(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 602, in read
    output = self._buffers[self._next_read_index].read(timeout)
  File "/home/ray/anaconda3/lib/python3.9/site-packages/ray/experimental/channel/shared_memory_channel.py", line 485, in read
    ret = self._worker.get_objects(
  File "python/ray/includes/common.pxi", line 100, in ray._raylet.check_status
ray.exceptions.RayChannelError: System error: Channel closed.
```

### Versions / Dependencies

```
>>> ray.__version__
'3.0.0.dev0'
>>> ray.__commit__
'62d59ba60823ceb7b2fb9a121c40ff58042be166'
```

### Reproduction script

https://github.com/anmscale/verl/pull/2

### Issue Severity

High: It blocks me from completing my task.


Hints:
I was able to reproduce the bug in the following standalone script. Note that it works fine when `nccl=False`.
https://github.com/anmscale/verl/blob/ray_cg/examples/ray/rcg_bug.py

For context, here's the workload I am trying to achieve:
![image](https://github.com/user-attachments/assets/98415937-0d13-4cc6-b517-157e15809cfc)

I think the issue is this part:

```
def scatter_from_rank_zero(
    workers: List[Worker], data: DataProto, nccl: bool = False
) -> Tuple[DataProto]:
    n_shards = len(workers)
    sharded_data = workers[0].chunk.options(num_returns=n_shards).bind(data, n_shards)
    type_hinted_data = [sharded_data[0]]
    for data in sharded_data[1:]:
        type_hinted_data.append(
            data.with_type_hint(TorchTensorType(transport="nccl")) if nccl else data
        )
    return tuple(type_hinted_data)
```

The type hint put on partial returns (e.g., when num_returns=n) is not currently respected. I will look into a fix.
Debugged further, the issue is missing data dependency when `num_returns` is used.

```(Worker pid=2902899) AssertionError: Channel cc21f308-88ff-4869-ac22-30130fc235aa does not exist in the buffer. ```
![acompiled_graph](https://github.com/user-attachments/assets/5bbc7146-81c6-4f7b-be93-08a432ecd586)

![acompiled_graph_schedule](https://github.com/user-attachments/assets/288b287f-9695-430f-8d7c-9ee5742e3350)


Patch:
diff --git a/python/ray/dag/dag_node_operation.py b/python/ray/dag/dag_node_operation.py
index ea71a88fb5e7..6073e58c72f7 100644
--- a/python/ray/dag/dag_node_operation.py
+++ b/python/ray/dag/dag_node_operation.py
@@ -435,7 +435,8 @@ def _build_dag_node_operation_graph(
             isinstance(task.dag_node, ClassMethodNode)
             and task.dag_node.is_class_method_output
         ):
-            # TODO(wxdeng): Handle the case where the task is a class method output.
+            # Class method output node dependencies are handled at its upstream:
+            # i.e., class method node
             continue
         for downstream_task_idx in task.downstream_task_idxs:
             downstream_dag_node = idx_to_task[downstream_task_idx].dag_node
@@ -445,8 +446,18 @@ def _build_dag_node_operation_graph(
                 isinstance(downstream_dag_node, ClassMethodNode)
                 and downstream_dag_node.is_class_method_output
             ):
-                # TODO(wxdeng): Handle the case where the downstream task is
-                # a class method output.
+                consumer_idxs = idx_to_task[downstream_task_idx].downstream_task_idxs
+                for consumer_idx in consumer_idxs:
+                    if consumer_idx in graph:
+                        _add_edge(
+                            graph[task_idx][_DAGNodeOperationType.WRITE],
+                            graph[consumer_idx][_DAGNodeOperationType.READ],
+                            "nccl"
+                            if graph[task_idx][
+                                _DAGNodeOperationType.WRITE
+                            ].requires_nccl
+                            else "shm",
+                        )
                 continue
             _add_edge(
                 graph[task_idx][_DAGNodeOperationType.WRITE],


Test patch:
diff --git a/python/ray/dag/tests/experimental/test_accelerated_dag.py b/python/ray/dag/tests/experimental/test_accelerated_dag.py
index bda5ecfb41eb..a9f1b3955783 100644
--- a/python/ray/dag/tests/experimental/test_accelerated_dag.py
+++ b/python/ray/dag/tests/experimental/test_accelerated_dag.py
@@ -223,6 +223,7 @@ def test_inc_two_returns(ray_start_regular, single_fetch):
         dag = MultiOutputNode([o1, o2])
 
     compiled_dag = dag.experimental_compile()
+    compiled_dag.visualize(channel_details=True)
     for i in range(3):
         refs = compiled_dag.execute(1)
         if single_fetch: