JSON encoding bug

Summary: This is a fix for #76 -- basically there were two separate issues, but both had to do with JSON encoding not working properly.

Reviewed By: kkashin

Differential Revision: D15314286

fbshipit-source-id: 92bafd5d462562d1fa671992cba72133155dd0a2

ax/storage/json_store/decoder.py

@@ -52,7 +52,9 @@ def object_from_json(object_json: Any) -> Any:
                 [(k, object_from_json(v)) for k, v in object_json["value"]]
             )
         elif _type == "DataFrame":
-            return pd.read_json(object_json["value"])
+            # Need dtype=False, otherwise infers arm_names like "4_1"
+            # should be int 41
+            return pd.read_json(object_json["value"], dtype=False)
         elif _type not in DECODER_REGISTRY:
             err = (
                 f"The JSON dictionary passed to `object_from_json` has a type "
@@ -210,6 +212,12 @@ def simple_experiment_from_json(object_json: Dict[str, Any]) -> SimpleExperiment
     experiment.is_test = object_from_json(is_test_json)
     experiment._time_created = object_from_json(time_created_json)
     experiment._trials = trials_from_json(experiment, trials_json)
+    for trial in experiment._trials.values():
+        for arm in trial.arms:
+            experiment._arms_by_signature[arm.signature] = arm
+    if experiment.status_quo is not None:
+        sq_sig = experiment.status_quo.signature
+        experiment._arms_by_signature[sq_sig] = experiment.status_quo
     experiment._experiment_type = object_from_json(experiment_type_json)
     experiment._data_by_trial = data_from_json(data_by_trial_json)
     return experiment

ax/storage/json_store/tests/test_json_store.py

@@ -34,7 +34,7 @@
     get_parameter_constraint,
     get_range_parameter,
     get_search_space,
-    get_simple_experiment,
+    get_simple_experiment_with_batch_trial,
     get_sum_constraint1,
     get_sum_constraint2,
     get_synthetic_runner,
@@ -61,7 +61,7 @@
     ("ParameterConstraint", get_parameter_constraint),
     ("RangeParameter", get_range_parameter),
     ("SearchSpace", get_search_space),
-    ("SimpleExperiment", get_simple_experiment),
+    ("SimpleExperiment", get_simple_experiment_with_batch_trial),
     ("SumConstraint", get_sum_constraint1),
     ("SumConstraint", get_sum_constraint2),
     ("SyntheticRunner", get_synthetic_runner),

ax/storage/sqa_store/tests/utils.py

@@ -26,7 +26,7 @@
     get_outcome_constraint,
     get_parameter_constraint,
     get_range_parameter,
-    get_simple_experiment,
+    get_simple_experiment_with_batch_trial,
     get_sum_constraint1,
     get_sum_constraint2,
     get_synthetic_runner,
@@ -138,7 +138,7 @@
     ),
     (
         "SimpleExperiment",
-        get_simple_experiment,
+        get_simple_experiment_with_batch_trial,
         Encoder.experiment_to_sqa,
         Decoder.experiment_from_sqa,
     ),

ax/utils/testing/fake.py

@@ -159,14 +159,25 @@ def get_factorial_experiment(
 
 def get_simple_experiment() -> SimpleExperiment:
     experiment = SimpleExperiment(
-        name="test_branin", search_space=get_branin_search_space(), objective_name="sum"
+        name="test_branin",
+        search_space=get_branin_search_space(),
+        status_quo=Arm(parameters={"x1": 0.0, "x2": 0.0}),
+        objective_name="sum",
     )
 
     experiment.description = "foobar"
 
     return experiment
 
 
+def get_simple_experiment_with_batch_trial() -> SimpleExperiment:
+    experiment = get_simple_experiment()
+    generator = get_sobol(experiment.search_space)
+    generator_run = generator.gen(10)
+    experiment.new_batch_trial(generator_run=generator_run)
+    return experiment
+
+
 def get_experiment_with_batch_trial() -> Experiment:
     batch_trial = get_batch_trial()
     return batch_trial.experiment