code: def setup_bench(
      platform, algo, bench_step, dataset, setup_kwargs, training_kwargs
  ):
      """
      Will setup the AlgorithmPair and the model to be ready for benchmark
  
      Parameters
      ----------
      platform :
         Either 'cpu' or 'cuml'
      algo_name :
         Algorithm/model name, can be found in the algorithms.py file
      bench_step :
          Either 'training' or 'inference', describe the algorithm/model
          step to be benchmarked
      dataset :
          Dataset data
      setup_kwargs :
          Algorithm/model setup kwargs
      training_kwargs :
          Algorithm/model training kwargs
      """
  
      # Generate the model
      if platform == "cuml":
          setup = algo.setup_cuml(dataset, **setup_kwargs)
      elif platform == "cpu":
          setup = algo.setup_cpu(dataset, **setup_kwargs)
  
      # Set the bench_func to perform training
      if bench_step == "training":
          if hasattr(algo.cuml_class, "fit"):
              algo.bench_func = fit
          # Model cannot be trained (special construction)
          elif algo.setup_cuml_func:
              pytest.skip(
                  "Model cannot be trained (special construction)"
              )
          else:
              raise ValueError("Training function not found")
      # Train the model and then set the bench_func to perform inference
      elif bench_step == "inference":
          if hasattr(algo.cuml_class, "fit"):
              algo.bench_func = fit
          # Model cannot be trained (special construction)
          elif algo.setup_cuml_func:
              algo.bench_func = pass_func
          else:
              raise ValueError("Training function not found")
  
          if platform == "cuml":
              setup["cuml_setup_result"] = algo.run_cuml(
                  dataset, bench_args=training_kwargs, **setup
              )
          elif platform == "cpu":
              setup["cpu_setup_result"] = algo.run_cpu(
                  dataset, bench_args=training_kwargs, **setup
              )
  
          if hasattr(algo.cuml_class, "predict"):
              algo.bench_func = predict
          elif hasattr(algo.cuml_class, "transform"):
              algo.bench_func = transform
          elif hasattr(algo.cuml_class, "kneighbors"):
              algo.bench_func = kneighbors
          elif any(
              hasattr(algo.cuml_class, attr)
              for attr in [
                  "fit_predict",
                  "fit_transform",
                  "fit_kneighbors",
              ]
          ):
              warnings.warn(
                  "Inference cannot be done separately, "
                  "doing both training and inference"
              )
              if hasattr(algo.cuml_class, "fit_predict"):
                  algo.bench_func = fit_predict
              elif hasattr(algo.cuml_class, "fit_transform"):
                  algo.bench_func = fit_transform
              elif hasattr(algo.cuml_class, "fit_kneighbors"):
                  algo.bench_func = fit_kneighbors
          else:
              raise ValueError("Inference function not found")
      else:
          raise ValueError(
              "bench_func should be either training or inference"
          )
      return setup
  
imports: import os
  import json
  import time
  import math
  import itertools
  import warnings
  from cuml.internals.safe_imports import cpu_only_import, gpu_only_import
  import pytest
  from cuml.benchmark import datagen, algorithms
  from cuml.benchmark.nvtx_benchmark import Profiler
  from dask.distributed import wait
  import dask.array
  import dask.dataframe
  from copy import copy
  from cuml.benchmark.bench_helper_funcs import pass_func, fit, predict, transform, kneighbors, fit_predict, fit_transform, fit_kneighbors
label_desc: Use image node_id 4 for calling a global function with example usage: setup_bench(platform, algo, bench_step, dataset, setup_kwargs, training_kwargs) and returns: setup