code: def main(args):
      ray.init()
      ray.data.DataContext.get_current().execution_options.verbose_progress = (
          True
      )
  
      dataset = ray.data.read_parquet(DATA_URI)
  
      if args.smoke_test:
          dataset = dataset.limit(1)
  
      actor_pool_size = int(ray.cluster_resources().get("GPU"))
      dataset = dataset.map_batches(
          GenerateImage,
          compute=ray.data.ActorPoolStrategy(size=actor_pool_size),
          batch_size=BATCH_SIZE,
          num_gpus=1,
      )
      dataset_iter = dataset.iter_batches(
          batch_format="pyarrow", batch_size=None
      )
  
      benchmark = Benchmark("stable_diffusion_benchmark")
      benchmark.run_iterate_ds("main", dataset_iter)
      test_output_json = os.environ.get(
          "TEST_OUTPUT_JSON", "release_test_out.json"
      )
      benchmark.write_result(test_output_json)
  
file_path: ray/release/nightly_tests/dataset/stable_diffusion_benchmark.py
imports: import argparse
  import os
  from typing import Dict
  import numpy
  import torch
  from diffusers import StableDiffusionImg2ImgPipeline
  from benchmark import Benchmark
  import ray
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types