How to properly save and load an experiment

I have a modified version of this https://botorch.org/tutorials/custom_botorch_model_in_ax

where I have saved the experiment after each call to get_botorch.

        for i in range(len(exp.trials.values()), num_bo_trails+2):
            print('Running optimization batch {}/{}'.format(i+1, num_bo_trails))
            model = get_botorch(experiment=exp, data=exp.eval(), search_space=exp.search_space,
                                model_constructor=_get_and_fit_gp)

            save(exp, args.bo_save_path)
            batch = exp.new_trial(generator_run=model.gen(1))

If that loop gets interupted, I want to be able to reload the experiment and restart the loop from where it left off. However I get his issue:

File “Torch1venv/venv/lib/python3.6/site-packages/ax/core/observation.py”, line 189, in observations_from_data obs_parameters = experiment.arms_by_name[features[“arm_name”]].parameters.copy() KeyError: ‘0’

After the first get_botorch call after I try to load up again.

Also I noticed that the trail status always seems to be ‘status=TrialStatus.RUNNING’ and never completed? Do I manually need to set trials to completed?

Thanks.

[ldworkin]

Hey @arvieFrydenlund – sorry, I forgot to mention this! It’s a simple fix on your end. After you load the experiment from the json file, you’ll just need to re-set the evaluation function, e.g.

exp = load(bo_save_path)
exp.evaluation_function = run

We don’t store evaluation functions, since function serialization is a difficult problem. We should make this more clear though 😃

[lena-kashtelyan]

Hey, @arvieFrydenlund, we tried to repro the bug you are getting, and coudn’t get the same issue to come up. Would you mind sharing your full notebook?