code: def main(
      num_runs: int = 1,
      num_trials: int = 8,
      num_workers: int = 4,
      use_gpu: bool = False,
      smoke_test: bool = False,
  ):
      ray.init(
          runtime_env={
              "working_dir": os.path.dirname(__file__),
              "env_vars": {"NCCL_SOCKET_IFNAME": "ens"},
          }
      )
      prepare_mnist()
  
      config = CONFIG.copy()
  
      if smoke_test:
          config["epochs"] = 1
  
      train_times = []
      tune_times = []
  
      for i in range(num_runs):
          print(f"Run {i+1} / {num_runs}")
  
          time.sleep(2)
  
          train_time = timeit.timeit(
              lambda: train_torch(
                  num_workers=num_workers,
                  use_gpu=use_gpu,
                  config=config,
              ),
              number=1,
          )
          train_times.append(train_time)
  
          time.sleep(2)
  
          tune_time = timeit.timeit(
              lambda: tune_torch(
                  num_workers=num_workers,
                  num_trials=num_trials,
                  use_gpu=use_gpu,
                  config=config,
              ),
              number=1,
          )
          tune_times.append(tune_time)
  
          result = {"train_time": train_time, "tune_time": tune_time}
          print(f"Results run {i+1}: {result}")
  
      mean_train_time = np.mean(train_times)
      mean_tune_time = np.mean(tune_times)
  
      full_results = {
          "train_times": train_times,
          "train_mean": mean_train_time,
          "train_sd": np.std(train_times),
          "tune_times": tune_times,
          "tune_mean": mean_tune_time,
          "tune_sd": np.std(tune_times),
      }
  
      print("Full results:", full_results)
  
      # NOTE: The value of `factor` is mostly arbitrary. It was previously `1.2`, but
      # that value turned out to be too low. For more context, see #29682.
      factor = 1.35
      threshold = mean_train_time * factor
  
      test_output_json = os.environ.get(
          "TEST_OUTPUT_JSON", "/tmp/result.json"
      )
      with open(test_output_json, "wt") as f:
          json.dump(full_results, f)
  
      assert (
          mean_tune_time <= threshold
      ), f"{mean_tune_time:.2f} > {threshold:.2f} = {factor:.1f} * {mean_train_time:.2f}"
  
file_path: ray/release/air_tests/air_benchmarks/workloads/tune_torch_benchmark.py
imports: import json
  import os
  import time
  import timeit
  from typing import Optional, Dict
  import click
  import numpy
  import ray
  from ray.train import ScalingConfig
  from ray.train.torch import TorchTrainer
label_desc: Use image node_id 5 for calling a global function with example usage: main(num_runs, num_trials, num_workers, use_gpu, smoke_test) without return types