bugfixes, set metric to allways be passthrough, override max_parallelism

boa/async-opt.py

@@ -9,11 +9,11 @@
 from ax import Data
 from ax.storage.json_store.decoder import object_from_json
 
-from boa.config import BOAConfig, BOAScriptOptions
+from boa.config import BOAConfig, BOAScriptOptions, MetricType
 from boa.controller import Controller
 from boa.storage import scheduler_from_json_file
 from boa.wrappers.synthetic_wrapper import SyntheticWrapper
-from boa.wrappers.wrapper_utils import cd_and_cd_back, load_jsonlike
+from boa.wrappers.wrapper_utils import load_jsonlike
 
 
 @click.command()
@@ -73,74 +73,70 @@ def main(config_path, scheduler_path, num_trials):
         if not config:
             sch_jsn = load_jsonlike(scheduler_path)
             config = BOAConfig(**{**object_from_json(sch_jsn["wrapper"]["config"]), **config_kw})
+    if "steps" in config.generation_strategy:
+        for step in config.generation_strategy["steps"]:
+            step.max_parallelism = None
+    else:
+        config.generation_strategy["max_parallelism_override"] = -1
+    for metric in config.objective.metrics:
+        metric.metric = "passthrough"
+        metric.metric_type = MetricType.PASSTHROUGH
 
-    options = dict(
-        append_timestamp=config.script_options.append_timestamp,
-        working_dir=config.script_options.working_dir,
-        wrapper=SyntheticWrapper(config=config),
-    )
-
-    with cd_and_cd_back(options["working_dir"]):
-        if scheduler_path:
-            scheduler = scheduler_from_json_file(filepath=scheduler_path)
-            if num_trials:
-                scheduler.wrapper.config.scheduler = dataclasses.replace(
-                    scheduler.wrapper.config.scheduler, total_trials=num_trials
-                )
-                scheduler.wrapper.config.n_trials = num_trials
-                scheduler.options = dataclasses.replace(scheduler.options, total_trials=num_trials)
-
-        else:
-            controller = Controller(config_path=config_path, wrapper=SyntheticWrapper(config=config))
-            controller.initialize_scheduler()
-            scheduler = controller.scheduler
-
-        if not scheduler.opt_csv.exists() and scheduler.experiment.trials:
-            controller.logger.warning(
-                "No optimization CSV found, but previous trials exist. "
-                "\nLikely cause was a previous run was moved with out the CSV."
+    if scheduler_path:
+        scheduler = scheduler_from_json_file(filepath=scheduler_path)
+        if num_trials:
+            scheduler.wrapper.config.scheduler = dataclasses.replace(
+                scheduler.wrapper.config.scheduler, total_trials=num_trials
             )
+            scheduler.wrapper.config.n_trials = num_trials
+            scheduler.options = dataclasses.replace(scheduler.options, total_trials=num_trials)
+
+    else:
+        controller = Controller(config_path=config_path, wrapper=SyntheticWrapper(config=config))
+        controller.initialize_scheduler()
+        scheduler = controller.scheduler
+
+    if not scheduler.opt_csv.exists() and scheduler.experiment.trials:
+        controller.logger.warning(
+            "No optimization CSV found, but previous trials exist. "
+            "\nLikely cause was a previous run was moved with out the CSV."
+        )
 
-        if scheduler.opt_csv.exists():
-            exp_attach_data_from_opt_csv(list(config.objective.metric_names), scheduler)
-            # for trial_index in new_data["trial_index"].unique():
-            #
-            #     scheduler.experiment.trials[trial_index].mark_completed()
+    if scheduler.opt_csv.exists():
+        exp_attach_data_from_opt_csv(list(config.objective.metric_names), scheduler)
 
-        generator_runs = [
-            scheduler.generation_strategy.gen(experiment=scheduler.experiment)
-            for _ in range(scheduler.wrapper.config.trials)
-        ]
+    generator_runs = scheduler.generation_strategy._gen_multiple(
+        experiment=scheduler.experiment, num_generator_runs=scheduler.wrapper.config.trials
+    )
 
-        for generator_run in generator_runs:
-            trial = scheduler.experiment.new_trial(
-                generator_run=generator_run,
-            )
-            trial.runner = scheduler.runner
-            trial.mark_running()
-
-        if scheduler.experiment.fetch_data().df.empty:
-            trials = scheduler.experiment.trials
-            metrics = scheduler.experiment.metrics
-            for metric in metrics.keys():
-                scheduler.experiment.attach_data(
-                    Data(
-                        df=pd.DataFrame.from_records(
-                            dict(
-                                trial_index=list(trials.keys()),
-                                arm_name=[f"{i}_0" for i in trials.keys()],
-                                metric_name=metric,
-                                mean=None,
-                                sem=0.0,
-                            )
+    for generator_run in generator_runs:
+        trial = scheduler.experiment.new_trial(
+            generator_run=generator_run,
+        )
+        trial.runner = scheduler.runner
+        trial.mark_running()
+
+    if scheduler.experiment.fetch_data().df.empty:
+        trials = scheduler.experiment.trials
+        metrics = scheduler.experiment.metrics
+        for metric in metrics.keys():
+            scheduler.experiment.attach_data(
+                Data(
+                    df=pd.DataFrame.from_records(
+                        dict(
+                            trial_index=list(trials.keys()),
+                            arm_name=[f"{i}_0" for i in trials.keys()],
+                            metric_name=metric,
+                            mean=None,
+                            sem=0.0,
                         )
-                    ),
-                    combine_with_last_data=True,
-                )
+                    )
+                ),
+                combine_with_last_data=True,
+            )
 
-        scheduler.save_data(metrics_to_end=True, ax_kwargs=dict(always_include_field_columns=True))
-        # scheduler =
-        return scheduler
+    scheduler.save_data(metrics_to_end=True, ax_kwargs=dict(always_include_field_columns=True))
+    return scheduler
 
 
 def exp_attach_data_from_opt_csv(metric_names, scheduler):
@@ -153,6 +149,8 @@ def exp_attach_data_from_opt_csv(metric_names, scheduler):
     nan_rows = exp_df["mean"].isna()
     nan_trials = exp_df.loc[nan_rows]["trial_index"].unique()
     new_data = df.loc[df["trial_index"].isin(nan_trials)]
+    if new_data.empty:
+        return
     for metric, trial_results in new_data[metric_names].to_dict().items():
         scheduler.experiment.attach_data(
             Data(

boa/storage.py

@@ -259,15 +259,7 @@ def scheduler_from_json_snapshot(
     )
 
     scheduler = Scheduler(generation_strategy=generation_strategy, experiment=experiment, options=options, **kwargs)
-    fp = serialized.pop("scheduler_filepath", None)
-    if fp:
-        scheduler.scheduler_filepath = object_from_json(
-            fp,
-            decoder_registry=decoder_registry,
-            class_decoder_registry=class_decoder_registry,
-        )
-    else:
-        scheduler.scheduler_filepath = pathlib.Path(filepath)
+    scheduler.scheduler_filepath = pathlib.Path(filepath)
     if wrapper:
         if isinstance(scheduler.experiment.runner, WrappedJobRunner):
             scheduler.experiment.runner.wrapper = wrapper