code: def run(
      num_runs: int = 1,
      num_epochs: int = 4,
      num_workers: int = 4,
      cpus_per_worker: int = 8,
      use_gpu: bool = False,
      batch_size: int = 64,
      smoke_test: bool = False,
      local: bool = False,
  ):
      # Note: smoke_test is ignored as we just adjust the batch size.
      # The parameter is passed by the release test pipeline.
      import ray
      from benchmark_util import (
          upload_file_to_all_nodes,
          run_command_on_all_nodes,
      )
  
      config = CONFIG.copy()
      config["epochs"] = num_epochs
      config["batch_size"] = batch_size
  
      if local:
          ray.init(num_cpus=4)
      else:
          ray.init("auto")
  
      print("Preparing Tensorflow benchmark: Downloading MNIST")
  
      path = str(
          (Path(__file__).parent / "_tensorflow_prepare.py").absolute()
      )
  
      upload_file_to_all_nodes(path)
      run_command_on_all_nodes(["python", path])
  
      times_ray = []
      times_local_ray = []
      losses_ray = []
      times_vanilla = []
      times_local_vanilla = []
      losses_vanilla = []
      for run in range(1, num_runs + 1):
          time.sleep(2)
  
          print(
              f"[Run {run}/{num_runs}] Running Tensorflow Ray benchmark"
          )
  
          time_ray, time_local_ray, loss_ray = train_tf_ray_air(
              num_workers=num_workers,
              cpus_per_worker=cpus_per_worker,
              use_gpu=use_gpu,
              config=config,
          )
  
          print(
              f"[Run {run}/{num_runs}] Finished Ray training ({num_epochs} epochs) in "
              f"{time_ray:.2f} seconds (local training time: {time_local_ray:.2f}s). "
              f"Observed loss = {loss_ray:.4f}"
          )
  
          time.sleep(2)
  
          print(
              f"[Run {run}/{num_runs}] Running Tensorflow vanilla benchmark"
          )
  
          # Todo: Vanilla runs are sometimes failing. We just retry here, but we should
          # get to the bottom of it.
          time_vanilla = time_local_vanilla = loss_vanilla = 0.0
          for i in range(3):
              try:
                  (
                      time_vanilla,
                      time_local_vanilla,
                      loss_vanilla,
                  ) = train_tf_vanilla(
                      num_workers=num_workers,
                      cpus_per_worker=cpus_per_worker,
                      use_gpu=use_gpu,
                      config=config,
                  )
              except Exception as e:
                  if i >= 2:
                      raise RuntimeError(
                          "Vanilla TF run failed 3 times"
                      ) from e
                  print("Vanilla TF run failed:", e)
                  continue
              break
  
          print(
              f"[Run {run}/{num_runs}] Finished vanilla training ({num_epochs} epochs) "
              f"in {time_vanilla:.2f} seconds "
              f"(local training time: {time_local_vanilla:.2f}s). "
              f"Observed loss = {loss_vanilla:.4f}"
          )
  
          print(
              f"[Run {run}/{num_runs}] Observed results: ",
              {
                  "tensorflow_mnist_ray_time_s": time_ray,
                  "tensorflow_mnist_ray_local_time_s": time_local_ray,
                  "tensorflow_mnist_ray_loss": loss_ray,
                  "tensorflow_mnist_vanilla_time_s": time_vanilla,
                  "tensorflow_mnist_vanilla_local_time_s": time_local_vanilla,
                  "tensorflow_mnist_vanilla_loss": loss_vanilla,
              },
          )
  
          times_ray.append(time_ray)
          times_local_ray.append(time_local_ray)
          losses_ray.append(loss_ray)
  
          times_vanilla.append(time_vanilla)
          times_local_vanilla.append(time_local_vanilla)
          losses_vanilla.append(loss_vanilla)
  
      times_ray_mean = np.mean(times_ray)
      times_ray_sd = np.std(times_ray)
  
      times_local_ray_mean = np.mean(times_local_ray)
      times_local_ray_sd = np.std(times_local_ray)
  
      times_vanilla_mean = np.mean(times_vanilla)
      times_vanilla_sd = np.std(times_vanilla)
  
      times_local_vanilla_mean = np.mean(times_local_vanilla)
      times_local_vanilla_sd = np.std(times_local_vanilla)
  
      result = {
          "tensorflow_mnist_ray_num_runs": num_runs,
          "tensorflow_mnist_ray_time_s_all": times_ray,
          "tensorflow_mnist_ray_time_s_mean": times_ray_mean,
          "tensorflow_mnist_ray_time_s_sd": times_ray_sd,
          "tensorflow_mnist_ray_time_local_s_all": times_local_ray,
          "tensorflow_mnist_ray_time_local_s_mean": times_local_ray_mean,
          "tensorflow_mnist_ray_time_local_s_sd": times_local_ray_sd,
          "tensorflow_mnist_ray_loss_mean": np.mean(losses_ray),
          "tensorflow_mnist_ray_loss_sd": np.std(losses_ray),
          "tensorflow_mnist_vanilla_time_s_all": times_vanilla,
          "tensorflow_mnist_vanilla_time_s_mean": times_vanilla_mean,
          "tensorflow_mnist_vanilla_time_s_sd": times_vanilla_sd,
          "tensorflow_mnist_vanilla_local_time_s_all": times_local_vanilla,
          "tensorflow_mnist_vanilla_local_time_s_mean": times_local_vanilla_mean,
          "tensorflow_mnist_vanilla_local_time_s_sd": times_local_vanilla_sd,
          "tensorflow_mnist_vanilla_loss_mean": np.mean(losses_vanilla),
          "tensorflow_mnist_vanilla_loss_std": np.std(losses_vanilla),
      }
  
      print("Results:", result)
      test_output_json = os.environ.get(
          "TEST_OUTPUT_JSON", "/tmp/result.json"
      )
      with open(test_output_json, "wt") as f:
          json.dump(result, f)
  
      target_ratio = 1.2
      ratio = (
          (times_local_ray_mean / times_local_vanilla_mean)
          if times_local_vanilla_mean != 0.0
          else 1.0
      )
      if ratio > target_ratio:
          raise RuntimeError(
              f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
              f"which is more than {target_ratio:.2f}x of the average vanilla training "
              f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). FAILED"
          )
  
      print(
          f"Training on Ray took an average of {times_local_ray_mean:.2f} seconds, "
          f"which is less than {target_ratio:.2f}x of the average vanilla training "
          f"time of {times_local_vanilla_mean:.2f} seconds ({ratio:.2f}x). PASSED"
      )
  
file_path: ray/release/air_tests/air_benchmarks/workloads/tensorflow_benchmark.py
imports: import json
  import os
  import time
  from pathlib import Path
  import click
  import numpy
  import tensorflow
  from typing import List, Tuple
label_desc: Use image node_id 8 for calling a global function with example usage: run(num_runs, num_epochs, num_workers, cpus_per_worker, use_gpu, batch_size, smoke_test, local) without return types