Runner in Developer API

Hi,

I’m using Ax stable (0.2.4) and trying to implement my own runner using the developer API. I follow this to write my runner. The code looks like this

from ax import Runner, Metric, Data, Experiment
from ax.modelbridge.registry import Models
from ax import OptimizationConfig, Objective

def ext_callable(params):
    return

class BirdRunner(Runner):
    def __init__(self, extra_param):
        self.more_arg = extra_param
    
    def run(self, trial):
        print('Runner is called, received trial:\n',trial)
        trial_metadata = {"name": str(trial.index)}
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            outp = ext_callable(params)
            trial_metadata.update(outp)

        return trial_metadata

class BirdMetric(Metric):
    def fetch_trial_data(self, trial):  
        records = []
        print('BirdMetric is called, fetching data from trial', trial)
        print(trial.deployed_name)
        return Data(df=pd.DataFrame.from_records(records))

optimization_config = OptimizationConfig(
    objective = Objective(
        metric=BirdMetric("my_bird_metric"), 
        minimize=True,
    ),
    outcome_constraints=[
    ],
)

search_space = SearchSpace(
    parameters=[
        RangeParameter(
            name="x1", parameter_type=ParameterType.FLOAT, lower=-2*np.pi, upper=2*np.pi
        ),
        RangeParameter(
            name="x2", parameter_type=ParameterType.FLOAT, lower=-2*np.pi, upper=2*np.pi
        ),
    ]
)

exp = Experiment(
    name="test_bird",
    search_space=search_space,
    optimization_config=optimization_config,
    runner=BirdRunner(extra_param={'arg1':0}),
)

NUM_SOBOL_TRIALS = 8

sobol_mdlbridge = Models.SOBOL(exp.search_space,seed=999)
batch_trials = exp.new_batch_trial(generator_run=sobol_mdlbridge.gen(NUM_SOBOL_TRIALS))
batch_trials.run()
batch_trials.mark_completed()

But I received an error. Here is the full traceback

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/utils/common/serialization.py in serialize_init_args(object, exclude_fields)
     75         try:
---> 76             value = getattr(object, arg)
     77         except AttributeError:

AttributeError: 'BirdRunner' object has no attribute 'extra_param'

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_50473/4273056987.py in <module>
     27 sobol_mdlbridge = Models.SOBOL(exp.search_space,seed=999)
     28 batch_trials = exp.new_batch_trial(generator_run=sobol_mdlbridge.gen(NUM_SOBOL_TRIALS))
---> 29 batch_trials.run()
     30 batch_trials.mark_completed()
     31 print(exp.fetch_data().df)

/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/core/batch_trial.py in run(self)
    449 
    450     def run(self) -> BatchTrial:
--> 451         return checked_cast(BatchTrial, super().run())
    452 
    453     def normalized_arm_weights(

/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/core/base_trial.py in run(self)
    382 
    383         # Default to experiment runner if trial doesn't have one
--> 384         self.assign_runner()
    385 
    386         if self._runner is None:

/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/core/base_trial.py in assign_runner(self)
    352         runner = self._runner or self.experiment.runner_for_trial(self)
    353         if runner is not None:
--> 354             self._runner = runner.clone()
    355         return self
    356 

/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/core/runner.py in clone(self)
    142         # pyre-ignore[45]: Cannot instantiate abstract class `Runner`.
    143         return cls(
--> 144             **serialize_init_args(self),
    145         )

/opt/anaconda3/envs/BO/lib/python3.7/site-packages/ax/utils/common/serialization.py in serialize_init_args(object, exclude_fields)
     77         except AttributeError:
     78             raise AttributeError(
---> 79                 f"{object.__class__} is missing a value for {arg}, "
     80                 f"which is needed by its constructor."
     81             )

AttributeError: <class '__main__.BirdRunner'> is missing a value for extra_param, which is needed by its constructor.

Can you help me understand how to write the runner properly? Thank you

[danielcohenlive]

What I’m seeing here is that in the first example, the metric name is injected as “bird_value” by the function bird(params). But under the hood we filter for the metrics we’re supposed to be fetching. So the first example prints a df in BirdMetric. fetch_trial_data() looking like:

fetch_trial_data   arm_name metric_name        mean  sem  trial_index
0      0_0  bird_value    4.031283  0.0            0
1      0_1  bird_value    4.246791  0.0            0
2      0_2  bird_value  118.204919  0.0            0
3      0_3  bird_value  152.538040  0.0            0
4      0_4  bird_value  -16.256667  0.0            0
5      0_5  bird_value    1.675551  0.0            0
6      0_6  bird_value   17.836805  0.0            0
7      0_7  bird_value   28.765350  0.0            0

while the second on prints:

  arm_name     metric_name        mean  sem  trial_index
0      0_0  my_bird_metric    4.031283  0.0            0
1      0_1  my_bird_metric    4.246791  0.0            0
2      0_2  my_bird_metric  118.204919  0.0            0
3      0_3  my_bird_metric  152.538040  0.0            0
4      0_4  my_bird_metric  -16.256667  0.0            0
5      0_5  my_bird_metric    1.675551  0.0            0
6      0_6  my_bird_metric   17.836805  0.0            0
7      0_7  my_bird_metric   28.765350  0.0            0

notice the difference in the metric_name column. The metric name you set on the opt config is “my_bird_metric”, so that’s why the results are filtered out in the first example.

Also, the runner shouldn’t really be setting data. The metric should. It should just do what the experiment needs to run in the environment it needs, but there might be special considerations I’m not aware of here.

[thgngu]

Thanks, I see my mistake now. Just to make it clearer for someone like myself who are not that proficient at programming

This won’t work:

class BirdRunner(Runner):
    def __init__(self, extra_param):
        self.more_arg = extra_param

this works:

class BirdRunner(Runner):
    def __init__(self, extra_param):
        self.extra_param= extra_param