code: def do_consume(
      split,
      rank,
      use_gpu=False,
      max_bytes_to_read=None,
  ):
      prefetch_batches = 1
      batch_size = 4096
      num_epochs = 1
  
      start = time.perf_counter()
      epochs_read, batches_read, bytes_read = 0, 0, 0
      batch_delays = []
  
      def generate_epochs(data, epochs: int):
          if isinstance(data, DatasetPipeline):
              for epoch in data.iter_epochs(epochs):
                  yield epoch
          elif isinstance(data, Dataset):
              # Dataset
              for _ in range(epochs):
                  yield data
          else:
              # StreamSplitDataIterator
              yield data
  
      for epoch_data in generate_epochs(split, num_epochs):
          epochs_read += 1
          batch_start = time.perf_counter()
  
          if isinstance(split, DatasetPipeline):
              batch_iterator = epoch_data.iter_batches(
                  prefetch_blocks=prefetch_batches,
                  batch_size=batch_size,
              )
          else:
              if not use_gpu:
                  batch_iterator = epoch_data.iter_batches(
                      prefetch_batches=prefetch_batches,
                      batch_size=batch_size,
                  )
              else:
                  batch_iterator = epoch_data.iter_torch_batches(
                      prefetch_batches=prefetch_batches,
                      batch_size=batch_size,
                      device="cuda",
                  )
  
          for batch in batch_iterator:
              batch_delay = time.perf_counter() - batch_start
              batch_delays.append(batch_delay)
              batches_read += 1
              if isinstance(batch, pd.DataFrame):
                  bytes_read += int(
                      batch.memory_usage(index=True, deep=True).sum()
                  )
              elif isinstance(batch, np.ndarray):
                  bytes_read += batch.nbytes
              elif isinstance(batch, dict) and isinstance(
                  batch.get("data"), torch.Tensor
              ):
                  tensor = batch["data"]
                  bytes_read += (
                      tensor.element_size() * tensor.nelement()
                  )
              else:
                  # NOTE: This isn't recursive and will just return the size of
                  # the object pointers if list of non-primitive types.
                  bytes_read += sys.getsizeof(batch)
              batch_start = time.perf_counter()
              if max_bytes_to_read is not None:
                  if bytes_read >= max_bytes_to_read:
                      break
      delta = time.perf_counter() - start
  
      print("Time to read all data", delta, "seconds")
      print(
          "P50/P95/Max batch delay (s)",
          np.quantile(batch_delays, 0.5),
          np.quantile(batch_delays, 0.95),
          np.max(batch_delays),
      )
      print("Num epochs read", epochs_read)
      print("Num batches read", batches_read)
      print(
          "Num bytes read", round(bytes_read / (1024 * 1024), 2), "MiB"
      )
      print(
          "Mean throughput",
          round(bytes_read / (1024 * 1024) / delta, 2),
          "MiB/s",
      )
  
      if rank == 0:
          print("Ingest stats from rank=0:\n\n{}".format(split.stats()))
  
file_path: ray/release/nightly_tests/dataset/data_ingest_benchmark.py
imports: import numpy
  import sys
  import time
  import argparse
  from benchmark import Benchmark
  import ray
  from ray.data import Dataset
  from ray.data import DatasetPipeline
  import pandas
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: do_consume(split, rank, use_gpu, max_bytes_to_read) without return types