code: def main(args):
      experiment = args.size if not args.smoke_test else "smoke_test"
      experiment_params = _EXPERIMENT_PARAMS[experiment]
  
      data_path, num_workers, cpus_per_worker = (
          experiment_params["data"],
          experiment_params["num_workers"],
          experiment_params["cpus_per_worker"],
      )
      print("Running xgboost training benchmark...")
      training_time = run_xgboost_training(
          data_path, num_workers, cpus_per_worker
      )
      print("Running xgboost prediction benchmark...")
      prediction_time = run_xgboost_prediction(
          _XGB_MODEL_PATH, data_path
      )
      result = {
          "training_time": training_time,
          "prediction_time": prediction_time,
      }
      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)
  
      if not args.disable_check:
          if training_time > _TRAINING_TIME_THRESHOLD:
              raise RuntimeError(
                  f"Training on XGBoost is taking {training_time} seconds, "
                  f"which is longer than expected ({_TRAINING_TIME_THRESHOLD} seconds)."
              )
  
          if prediction_time > _PREDICTION_TIME_THRESHOLD:
              raise RuntimeError(
                  f"Batch prediction on XGBoost is taking {prediction_time} seconds, "
                  f"which is longer than expected ({_PREDICTION_TIME_THRESHOLD} seconds)."
              )
  
file_path: ray/release/air_tests/air_benchmarks/workloads/xgboost_benchmark.py
imports: from functools import wraps
  import json
  import multiprocessing
  from multiprocessing import Process
  import numpy
  import os
  import pandas
  import time
  import traceback
  from typing import Dict
  import xgboost
  import ray
  from ray import data
  from ray.train.xgboost import XGBoostTrainer
  from ray.train import RunConfig, ScalingConfig
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types