code: def test_tensor_parallel_dag(shared_ray_instance):
      """Simulate the TP DAG with N workers.
      Input -> forward -> MultiOutput
      """
  
      @ray.remote
      class Worker:
          def __init__(self, rank):
              self.rank = rank
              self.forwarded = 0
  
          def forward(self, input_data: int):
              print(input_data)
              self.forwarded += 1
              return self.rank + input_data
  
          def initialize(self):
              pass
  
          def get_forwarded(self):
              return self.forwarded
  
      NUM_WORKERS = 4
      workers = [Worker.remote(i) for i in range(NUM_WORKERS)]
      # Init multiple times.
      for _ in range(4):
          ray.get([worker.initialize.remote() for worker in workers])
  
      with InputNode() as input_data:
          dag = MultiOutputNode(
              [worker.forward.bind(input_data) for worker in workers]
          )
  
      # Run DAG repetitively.
      ITER = 4
      assert ITER > 1
      for i in range(ITER):
          ref = dag.execute(i)
          all_outputs = ray.get(ref)
          assert len(all_outputs) == NUM_WORKERS
          assert all_outputs == [i + j for j in range(NUM_WORKERS)]
  
      forwarded = ray.get(
          [worker.get_forwarded.remote() for worker in workers]
      )
      assert forwarded == [ITER for _ in range(NUM_WORKERS)]
  
file_path: ray/python/ray/dag/tests/test_output_node.py
imports: import pytest
  import ray
  from ray.dag.input_node import InputNode
  from ray.dag.output_node import MultiOutputNode
  from ray.util.state import list_tasks
  from ray._private.test_utils import wait_for_condition
label_desc: Use image node_id 4 for calling a global function with example usage: test_tensor_parallel_dag(shared_ray_instance) and returns: unknown, self