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 Torch benchmark: Downloading MNIST")
  
      path = str(
          (Path(__file__).parent / "_torch_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 Torch Ray benchmark")
  
          time_ray, time_local_ray, loss_ray = train_torch_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 Torch vanilla benchmark"
          )
  
          (
              time_vanilla,
              time_local_vanilla,
              loss_vanilla,
          ) = train_torch_vanilla(
              num_workers=num_workers,
              cpus_per_worker=cpus_per_worker,
              use_gpu=use_gpu,
              config=config,
          )
  
          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 = {
          "torch_mnist_ray_num_runs": num_runs,
          "torch_mnist_ray_time_s_all": times_ray,
          "torch_mnist_ray_time_s_mean": times_ray_mean,
          "torch_mnist_ray_time_s_sd": times_ray_sd,
          "torch_mnist_ray_time_local_s_all": times_local_ray,
          "torch_mnist_ray_time_local_s_mean": times_local_ray_mean,
          "torch_mnist_ray_time_local_s_sd": times_local_ray_sd,
          "torch_mnist_ray_loss_mean": np.mean(losses_ray),
          "torch_mnist_ray_loss_sd": np.std(losses_ray),
          "torch_mnist_vanilla_time_s_all": times_vanilla,
          "torch_mnist_vanilla_time_s_mean": times_vanilla_mean,
          "torch_mnist_vanilla_time_s_sd": times_vanilla_sd,
          "torch_mnist_vanilla_local_time_s_all": times_local_vanilla,
          "torch_mnist_vanilla_local_time_s_mean": times_local_vanilla_mean,
          "torch_mnist_vanilla_local_time_s_sd": times_local_vanilla_sd,
          "torch_mnist_vanilla_loss_mean": np.mean(losses_vanilla),
          "torch_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.15
      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"
      )
  
imports: import json
  import os
  import time
  from pathlib import Path
  from typing import Dict, Tuple
  import click
  import numpy
  import torch
  from torch import nn, distributed
  from torch.utils.data import DataLoader, DistributedSampler
  from torch.utils.data.dataloader import default_collate
  from torchvision import datasets
  from torchvision.transforms import ToTensor
label_desc: Use image node_id 9 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