docstr: Runs optimization over k x n rounds of k parallel trials.
  
  This asks Ax to run up to max_parallelism_cap trials in parallel by
  submitting them to the scheduler at the same time.
  
  NOTE:
      * setting max_parallelism_cap in generation_strategy
      * setting run_trials_in_batches in scheduler options
      * setting total_trials = parallelism * rounds
docstr_tok: ["Runs","optimization","over","k","x","n","rounds","of","k","parallel","trials",".","This","asks","Ax","to","run","up","to","max_parallelism_cap","trials","in","parallel","by","submitting","them","to","the","scheduler","at","the","same","time",".","NOTE",":","*","setting","max_parallelism_cap","in","generation_strategy","*","setting","run_trials_in_batches","in","scheduler","options","*","setting","total_trials","=","parallelism","*","rounds"]
code: def test_run_experiment_locally_in_batches(self) -> None:
      """Runs optimization over k x n rounds of k parallel trials.
  
      This asks Ax to run up to max_parallelism_cap trials in parallel by
      submitting them to the scheduler at the same time.
  
      NOTE:
          * setting max_parallelism_cap in generation_strategy
          * setting run_trials_in_batches in scheduler options
          * setting total_trials = parallelism * rounds
  
      """
      parallelism = 2
      rounds = 3
  
      experiment = Experiment(
          name="torchx_booth_parallel_demo",
          search_space=SearchSpace(parameters=self._parameters),
          optimization_config=OptimizationConfig(
              objective=self._objective
          ),
          runner=self._runner,
          is_test=True,
          properties={Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF: True},
      )
  
      scheduler = Scheduler(
          experiment=experiment,
          generation_strategy=(
              choose_generation_strategy(
                  search_space=experiment.search_space,
                  max_parallelism_cap=parallelism,
              )
          ),
          options=SchedulerOptions(
              run_trials_in_batches=True,
              total_trials=(parallelism * rounds),
          ),
      )
  
      try:
          scheduler.run_all_trials()
  
          # TorchXMetric always returns trial index; hence the best experiment
          # for min objective will be the params for trial 0.
          scheduler.report_results()
      except FailureRateExceededError:
          pass
  
file_path: Ax/ax/runners/tests/test_torchx.py
imports: import os
  import shutil
  import tempfile
  from typing import List
  from ax.core import BatchTrial, Experiment, Objective, OptimizationConfig, Parameter, ParameterType, RangeParameter, SearchSpace
  from ax.metrics.torchx import TorchXMetric
  from ax.modelbridge.dispatch_utils import choose_generation_strategy
  from ax.runners.torchx import TorchXRunner
  from ax.service.scheduler import FailureRateExceededError, Scheduler, SchedulerOptions
  from ax.utils.common.constants import Keys
  from ax.utils.common.testutils import TestCase
  from torchx.components import utils
label_desc: Use image node_id 5 for calling the TorchXRunnerTest obj's underlying member method code with example usage: obj.test_run_experiment_locally_in_batches() without return types