class_docstr: Utility class used for Ray Data release tests and benchmarks, which works
  for both local and distributed benchmarking. When run on the nightly release
  test pipeline, the results are written to our internal database, which
  can then be rendered on the dashboard. Usage tips:
  
  A typical workflow would be:
  
  benchmark = Benchmark("benchmark-name")
  
  # set up (such as input read or generation)
  ...
  
  benchmark.run_materialize_ds("case-1", fn_1)
  # Could be Ray Data iterator, Torch DataLoader, TF Dataset...
  benchmark.run_iterate_ds("case-2", dataset)
  benchmark.run_fn("case-3", fn_3)
  
  # Writes a JSON with metrics of the form:
  # {"case-1": {...}, "case-2": {...}, "case-3": {...}}
  benchmark.write_result()
  
  See example usage in ``aggregate_benchmark.py``.
class_docstr_tok: ["Utility","class","used","for","Ray","Data","release","tests","and","benchmarks",",","which","works","for","both","local","and","distributed","benchmarking",".","When","run","on","the","nightly","release","test","pipeline",",","the","results","are","written","to","our","internal","database",",","which","can","then","be","rendered","on","the","dashboard",".","Usage","tips",":","A","typical","workflow","would","be",":","benchmark","=","Benchmark","(","``","benchmark-name","''",")","#","set","up","(","such","as","input","read","or","generation",")","...","benchmark.run_materialize_ds","(","``","case-1","''",",","fn_1",")","#","Could","be","Ray","Data","iterator",",","Torch","DataLoader",",","TF","Dataset","...","benchmark.run_iterate_ds","(","``","case-2","''",",","dataset",")","benchmark.run_fn","(","``","case-3","''",",","fn_3",")","#","Writes","a","JSON","with","metrics","of","the","form",":","#","{","``","case-1","''",":","{","...","}",",","``","case-2","''",":","{","...","}",",","``","case-3","''",":","{","...","}","}","benchmark.write_result","(",")","See","example","usage","in","``","aggregate_benchmark.py","``","."]
docstr: Run a benchmark iterating over a dataset. Runtime and throughput
  are automatically calculated and reported. Supported dataset types are:
  - Ray Dataset (`ray.data.Dataset`)
  - iterator over Ray Dataset (`ray.data.iterator._IterableFromIterator` from
      `.iter_batches()`,`.iter_torch_batches()`, `.iter_tf_batches()`)
  - Torch DataLoader (`torch.utils.data.DataLoader`)
  - TensorFlow Dataset (`tf.data.Dataset`)
docstr_tok: ["Run","a","benchmark","iterating","over","a","dataset",".","Runtime","and","throughput","are","automatically","calculated","and","reported",".","Supported","dataset","types","are",":","-","Ray","Dataset","(","`","ray.data.Dataset","`",")","-","iterator","over","Ray","Dataset","(","`","ray.data.iterator._IterableFromIterator","`","from","`",".iter_batches","(",")","`",",","`",".iter_torch_batches","(",")","`",",","`",".iter_tf_batches","(",")","`",")","-","Torch","DataLoader","(","`","torch.utils.data.DataLoader","`",")","-","TensorFlow","Dataset","(","`","tf.data.Dataset","`",")"]
code: def run_iterate_ds(
      self,
      name: str,
      dataset: Any,
  ):
      """Run a benchmark iterating over a dataset. Runtime and throughput
      are automatically calculated and reported. Supported dataset types are:
      - Ray Dataset (`ray.data.Dataset`)
      - iterator over Ray Dataset (`ray.data.iterator._IterableFromIterator` from
          `.iter_batches()`,`.iter_torch_batches()`, `.iter_tf_batches()`)
      - Torch DataLoader (`torch.utils.data.DataLoader`)
      - TensorFlow Dataset (`tf.data.Dataset`)
      """
      # Import TF/Torch within this method, as not all benchmarks
      # will use/install these libraries.
      import tensorflow as tf
      import torch
  
      gc.collect()
  
      print(f"Running case: {name}")
      start_time = time.perf_counter()
      record_count = 0
      ds_iterator = iter(dataset)
      for batch in ds_iterator:
          # Unwrap list to get the underlying batch format.
          if isinstance(batch, (list, tuple)) and len(batch) > 0:
              batch = batch[0]
  
          # Get the batch size for various batch formats.
          if isinstance(batch, dict):
              feature_lengths = {k: len(batch[k]) for k in batch}
              batch_size = max(feature_lengths.values())
          elif isinstance(batch, (pa.Table, pd.DataFrame)):
              batch_size = len(batch)
          elif isinstance(batch, torch.Tensor):
              batch_size = batch.size(dim=0)
          elif isinstance(batch, tf.Tensor):
              batch_size = batch.shape.as_list()[0]
          else:
              raise TypeError(f"Unexpected batch type: {type(batch)}")
          record_count += batch_size
  
      duration = time.perf_counter() - start_time
      self.result[name] = {
          BenchmarkMetric.RUNTIME.value: duration,
          BenchmarkMetric.NUM_ROWS.value: record_count,
          BenchmarkMetric.THROUGHPUT.value: record_count / duration,
      }
      print(f"Result of case {name}: {self.result[name]}")
  
file_path: ray/release/nightly_tests/dataset/benchmark.py
imports: import gc
  import json
  import os
  import time
  from typing import Callable, Dict
  from ray.data.dataset import Dataset
  from typing import Any
  from enum import Enum
  import pyarrow
  import pandas
label_desc: Use image node_id 3 for calling the Benchmark obj's underlying member method code with example usage: obj.run_iterate_ds(name, dataset) without return types