code: def run(
      ray_address=None,
      object_store_memory=1e9,
      num_partitions=5,
      partition_size=200e6,
      num_nodes=None,
      num_cpus=8,
      no_streaming=False,
      use_wait=False,
      tracker=None,
  ):
      import time
  
      import numpy as np
  
      is_multi_node = num_nodes
      if ray_address:
          print("Connecting to a existing cluster...")
          ray.init(address=ray_address, ignore_reinit_error=True)
      elif is_multi_node:
          print("Emulating a cluster...")
          print(f"Num nodes: {num_nodes}")
          print(f"Num CPU per node: {num_cpus}")
          print(f"Object store memory per node: {object_store_memory}")
          cluster = build_cluster(
              num_nodes, num_cpus, object_store_memory
          )
          ray.init(address=cluster.address)
      else:
          print("Start a new cluster...")
          ray.init(
              num_cpus=num_cpus, object_store_memory=object_store_memory
          )
  
      partition_size = int(partition_size)
      num_partitions = num_partitions
      rows_per_partition = partition_size // (8 * 2)
      if tracker is None:
          tracker = _StatusTracker.remote()
      use_wait = use_wait
  
      def input_reader(i: PartitionID) -> Iterable[InType]:
          for _ in range(num_partitions):
              yield np.ones(
                  (rows_per_partition // num_partitions, 2),
                  dtype=np.int64,
              )
  
      def output_writer(
          i: PartitionID, shuffle_inputs: List[ObjectRef]
      ) -> OutType:
          total = 0
          if not use_wait:
              for obj_ref in shuffle_inputs:
                  arr = ray.get(obj_ref)
                  total += arr.size * arr.itemsize
          else:
              while shuffle_inputs:
                  [ready], shuffle_inputs = ray.wait(
                      shuffle_inputs, num_returns=1
                  )
                  arr = ray.get(ready)
                  total += arr.size * arr.itemsize
  
          return total
  
      def output_writer_non_streaming(
          i: PartitionID, shuffle_inputs: List[Any]
      ) -> OutType:
          total = 0
          for arr in shuffle_inputs:
              total += arr.size * arr.itemsize
          return total
  
      if no_streaming:
          output_writer_callable = output_writer_non_streaming
          object_store_writer = ObjectStoreWriterNonStreaming
      else:
          object_store_writer = ObjectStoreWriter
          output_writer_callable = output_writer
  
      start = time.time()
      output_sizes = simple_shuffle(
          input_reader=input_reader,
          input_num_partitions=num_partitions,
          output_num_partitions=num_partitions,
          output_writer=output_writer_callable,
          object_store_writer=object_store_writer,
          tracker=tracker,
      )
      delta = time.time() - start
  
      time.sleep(0.5)
      print()
  
      summary = None
      for i in range(5):
          try:
              summary = ray._private.internal_api.memory_summary(
                  stats_only=True
              )
          except Exception:
              time.sleep(1)
              pass
          if summary:
              break
      print(summary)
      print()
      print(
          "Shuffled",
          int(sum(output_sizes) / (1024 * 1024)),
          "MiB in",
          delta,
          "seconds",
      )
  
imports: import time
  from typing import Any, Callable, Iterable, List, Tuple, Union
  import ray
  from ray import ObjectRef
  from ray.cluster_utils import Cluster
label_desc: Use image node_id 5 for calling a global function with example usage: run(ray_address, object_store_memory, num_partitions, partition_size, num_nodes, num_cpus, no_streaming, use_wait, tracker) and returns: total, total