Serialize _trial_type_to_metric_names and use base class for SQA encoding (facebook#4999)

Summary:

Phase 5 of moving MultiTypeExperiment features into base Experiment.

JSON storage:
- `experiment_to_dict` now serializes `default_trial_type` and `_trial_type_to_metric_names` (sets converted to sorted lists for JSON determinism).
- `experiment_from_json` pops and restores `_trial_type_to_metric_names` (lists back to sets), defaulting to None for backward compat with old JSON.
- `multi_type_experiment_from_json` pops `_trial_type_to_metric_names` (now present via `experiment_to_dict`) to prevent it being passed as a constructor kwarg.

SQA storage:
- Encoder replaces the `isinstance(experiment, MultiTypeExperiment)` main branch with `experiment._default_trial_type is not None`, using the base class `metric_to_trial_type` computed property instead of MTE-specific `_metric_to_trial_type`. Keeps `isinstance` only for MTE-specific fields (`Keys.SUBCLASS` marker, `_metric_to_canonical_name`).
- No decoder changes needed — `_trial_type_to_metric_names` is already populated correctly through the `__init__` + `add_tracking_metric` call chain.

Reviewed By: lena-kashtelyan

Differential Revision: D94994120

Author: mpolson64

ax/storage/json_store/decoder.py

@@ -702,6 +702,9 @@ def multi_type_experiment_from_json(
     )
     # not relevant to multi type experiment
     del object_json["runner"]
+    # Pop _trial_type_to_metric_names if present (not a constructor argument;
+    # rebuilt from _metric_to_trial_type below).
+    object_json.pop("_trial_type_to_metric_names", None)
 
     kwargs = {
         k: object_from_json(
@@ -752,6 +755,11 @@ def experiment_from_json(
         if _trial_type_to_runner_json is not None
         else None
     )
+    # Pop _trial_type_to_metric_names before constructing Experiment
+    # (it's not a constructor argument). Convert lists back to sets.
+    _trial_type_to_metric_names_json = object_json.pop(
+        "_trial_type_to_metric_names", None
+    )
 
     experiment = Experiment(
         **{
@@ -766,6 +774,10 @@ def experiment_from_json(
     experiment._arms_by_name = {}
     if _trial_type_to_runner is not None:
         experiment._trial_type_to_runner = _trial_type_to_runner
+    if _trial_type_to_metric_names_json is not None:
+        experiment._trial_type_to_metric_names = {
+            tt: set(names) for tt, names in _trial_type_to_metric_names_json.items()
+        }
 
     _load_experiment_info(
         exp=experiment,

ax/storage/json_store/encoders.py

@@ -128,17 +128,20 @@ def experiment_to_dict(experiment: Experiment) -> dict[str, Any]:
         "is_test": experiment.is_test,
         "data_by_trial": data_to_data_by_trial(data=experiment.data),
         "properties": experiment._properties,
+        "default_trial_type": experiment._default_trial_type,
         "_trial_type_to_runner": experiment._trial_type_to_runner,
+        "_trial_type_to_metric_names": {
+            tt: sorted(names)
+            for tt, names in experiment._trial_type_to_metric_names.items()
+        },
     }
 
 
 def multi_type_experiment_to_dict(experiment: MultiTypeExperiment) -> dict[str, Any]:
     """Convert AE multitype experiment to a dictionary."""
     multi_type_dict = {
-        "default_trial_type": experiment._default_trial_type,
         "_metric_to_canonical_name": experiment._metric_to_canonical_name,
         "_metric_to_trial_type": experiment._metric_to_trial_type,
-        "_trial_type_to_runner": experiment._trial_type_to_runner,
     }
     multi_type_dict.update(experiment_to_dict(experiment))
     return multi_type_dict

ax/storage/sqa_store/encoder.py

@@ -222,13 +222,19 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
         runners = []
         if isinstance(experiment, MultiTypeExperiment):
             experiment._properties[Keys.SUBCLASS] = "MultiTypeExperiment"
+        if experiment._default_trial_type is not None:
             for trial_type, runner in experiment._trial_type_to_runner.items():
                 runner_sqa = self.runner_to_sqa(none_throws(runner), trial_type)
                 runners.append(runner_sqa)
 
+            metric_to_tt = experiment.metric_to_trial_type
             for metric in tracking_metrics:
-                metric.trial_type = experiment._metric_to_trial_type[metric.name]
-                if metric.name in experiment._metric_to_canonical_name:
+                if metric.name in metric_to_tt:
+                    metric.trial_type = metric_to_tt[metric.name]
+                if (
+                    isinstance(experiment, MultiTypeExperiment)
+                    and metric.name in experiment._metric_to_canonical_name
+                ):
                     metric.canonical_name = experiment._metric_to_canonical_name[
                         metric.name
                     ]