code: def benchmark_code(
      args,
  ):
      cache_input_ds = args.cache_input_ds
      cache_output_ds = args.cache_output_ds
      assert (
          sum([cache_output_ds, cache_input_ds]) <= 1
      ), "Can only test one caching variant at a time"
  
      if args.use_torch or args.split_input:
          split_input_files_per_worker(args)
  
      ray_datasets_dict = {}
      if not (args.use_mosaic or args.use_torch):
          # Only create one dataset if `args.split_input` is True.
          # Otherwise, create N datasets for N training workers,
          # each dataset reads the corresponding portion of input data.
          num_datasets = 1
          if args.split_input:
              num_datasets = args.num_workers
  
          for i in range(num_datasets):
              if args.split_input:
                  input_paths = INPUT_FILES_PER_WORKER[i]
                  ds_name = f"train_{i}"
              else:
                  input_paths = args.data_root
                  ds_name = "train"
  
              # 1) Read in data with read_images() / read_parquet()
              if args.file_type == "image":
                  # Obtain WNID from filepath, then convert to numerical class ID
                  partitioning = Partitioning(
                      "dir",
                      field_names=["class"],
                      base_dir=args.data_root,
                  )
                  ray_dataset = ray.data.read_images(
                      input_paths,
                      mode="RGB",
                      shuffle="files",
                      partitioning=partitioning,
                  )
  
                  val_dataset = ray.data.Dataset.copy(ray_dataset)
                  # Full random shuffle results in OOM. Instead, use the
                  # `ds.iter_batches(local_shuffle_buffer_size=...)`
                  # parameter in the training loop.
                  # ray_dataset = ray_dataset.random_shuffle()
  
                  def wnid_to_index(row):
                      row["label"] = IMAGENET_WNID_TO_ID[row["class"]]
                      row.pop("class")
                      return row
  
                  ray_dataset = ray_dataset.map(wnid_to_index)
                  val_dataset = val_dataset.map(wnid_to_index)
              elif args.file_type == "parquet":
                  ray_dataset = ray.data.read_parquet(
                      args.data_root,
                  )
              else:
                  raise Exception(f"Unknown file type {args.file_type}")
              if cache_input_ds:
                  ray_dataset = ray_dataset.materialize()
  
              # 2) Preprocess data by applying transformation with map/map_batches()
              if args.file_type == "image":
                  ray_dataset = ray_dataset.map(crop_and_flip_image)
                  val_dataset = val_dataset.map(center_crop_image)
                  ray_datasets_dict["val"] = val_dataset
              elif args.file_type == "parquet":
                  ray_dataset = ray_dataset.map(
                      decode_image_crop_and_flip
                  )
              if cache_output_ds:
                  ray_dataset = ray_dataset.materialize()
              ray_datasets_dict[ds_name] = ray_dataset
  
      # 3) Train TorchTrainer on processed data
      options = DataConfig.default_ingest_options()
      options.preserve_order = args.preserve_order
  
      torch_trainer = TorchTrainer(
          train_loop_per_worker,
          datasets=ray_datasets_dict,
          scaling_config=ScalingConfig(
              num_workers=args.num_workers,
              use_gpu=args.use_gpu,
          ),
          dataset_config=ray.train.DataConfig(
              datasets_to_split=[] if args.split_input else "all",
              execution_options=options,
          ),
          run_config=RunConfig(
              storage_path="/mnt/cluster_storage",
              failure_config=train.FailureConfig(args.num_retries),
          ),
      )
  
      result = torch_trainer.fit()
      data_benchmark_metrics = {}
  
      # Report the average of per-epoch throughput, excluding the first epoch.
      # Unless there is only one epoch, in which case we report the epoch
      # throughput directly.
      start_epoch_tput = 0 if args.num_epochs == 1 else 1
      epoch_tputs = []
      num_rows_per_epoch = sum(result.metrics["num_rows"])
      for i in range(start_epoch_tput, args.num_epochs):
          time_start_epoch_i, time_end_epoch_i = zip(
              *result.metrics[f"epoch_{i}_times"]
          )
          runtime_epoch_i = max(time_end_epoch_i) - min(
              time_start_epoch_i
          )
          tput_epoch_i = num_rows_per_epoch / runtime_epoch_i
          epoch_tputs.append(tput_epoch_i)
      avg_per_epoch_tput = sum(epoch_tputs) / len(epoch_tputs)
      print(
          "Total num rows read per epoch:", num_rows_per_epoch, "images"
      )
      print(
          "Averaged per-epoch throughput:", avg_per_epoch_tput, "img/s"
      )
      data_benchmark_metrics.update(
          {
              BenchmarkMetric.THROUGHPUT: avg_per_epoch_tput,
          }
      )
  
      # Report the training accuracy of the final epoch.
      if result.metrics.get("num_rows_val") is not None:
          final_epoch_acc = sum(
              result.metrics["num_rows_correct_val"]
          ) / sum(result.metrics["num_rows_val"])
          print(f"Final epoch accuracy: {final_epoch_acc * 100:.3f}%")
          data_benchmark_metrics.update(
              {
                  BenchmarkMetric.ACCURACY: final_epoch_acc,
              }
          )
      return data_benchmark_metrics
  
imports: import ray
  from ray import train
  from ray.train import DataConfig, ScalingConfig, RunConfig, Checkpoint
  from ray.train.torch import TorchTrainer
  from ray.data.datasource.partitioning import Partitioning
  import tempfile
  import itertools
  import os
  import time
  from benchmark import Benchmark, BenchmarkMetric
  from image_loader_microbenchmark import get_transform, crop_and_flip_image, decode_image_crop_and_flip, center_crop_image
  from image_loader_microbenchmark import get_mosaic_dataloader
  import torch
  import torch.distributed
  from torchvision.models import resnet50
  from torchvision.transforms.functional import pil_to_tensor
  import torch.nn
  import torch.optim
  from dataset_benchmark_util import get_prop_parquet_paths, IMG_S3_ROOT, get_mosaic_epoch_size, IMAGENET_WNID_TO_ID
label_desc: Use image node_id 7 for calling a global function with example usage: benchmark_code(args) and returns: data_benchmark_metrics, row