Centralize experiment property serialization for JSON storage (#5113)

Summary:
Pull Request resolved: #5113

`update_properties_on_experiment` writes `experiment._properties` directly to the DB via raw `json.dumps`, bypassing the `experiment_to_sqa` encoder that converts non-JSON-serializable objects (like `LLMMessage` dataclasses) to plain dicts. This causes `TypeError: Object of type LLMMessage is not JSON serializable` for any experiment with LLM messages saved through the incremental property update path.

This diff extracts the property preparation logic into a shared `prepare_experiment_properties_for_storage` function called from both `experiment_to_sqa` and `update_properties_on_experiment`, ensuring all save paths use consistent serialization.

We considered adding a blanket `default=dataclasses.asdict` handler to the `JSONEncodedObject` TypeDecorator, but rejected it because: (1) dataclasses with non-JSON fields (e.g. `BenchmarkTrialMetadata` with `pd.DataFrame`) would produce confusing partial-conversion errors, (2) silently serializing dataclasses without paired decoders breaks round-trip fidelity, and (3) it removes the fail-fast guardrail for unexpected objects in JSON columns.

Differential Revision: D98828500

fbshipit-source-id: ed458af8647d30ac018e885e00a3466daeec6104

Author: ItsMrLin

ax/storage/sqa_store/encoder.py

@@ -88,6 +88,31 @@
 logger: Logger = get_logger(__name__)
 
 
+def prepare_experiment_properties_for_storage(
+    experiment: Experiment,
+) -> dict[str, Any]:
+    """Prepare experiment properties for JSON storage by converting non-JSON-
+    serializable objects (e.g. dataclasses) to plain dicts.
+
+    This is the single source of truth for experiment property serialization.
+    All code paths that persist experiment properties to the database should
+    use this function to ensure consistent handling.
+    """
+    properties = experiment._properties.copy()
+    if (
+        oc := experiment.optimization_config
+    ) is not None and oc.pruning_target_parameterization is not None:
+        properties["pruning_target_parameterization"] = arm_to_dict(
+            oc.pruning_target_parameterization
+        )
+    if Keys.LLM_MESSAGES in properties:
+        properties[Keys.LLM_MESSAGES] = [
+            dataclasses.asdict(m) if isinstance(m, LLMMessage) else m
+            for m in properties[Keys.LLM_MESSAGES]
+        ]
+    return properties
+
+
 class Encoder:
     """Class that contains methods for storing an Ax experiment to SQLAlchemy.
 
@@ -230,18 +255,7 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
                     ]
         elif experiment.runner:
             runners.append(self.runner_to_sqa(none_throws(experiment.runner)))
-        properties = experiment._properties.copy()
-        if (
-            oc := experiment.optimization_config
-        ) is not None and oc.pruning_target_parameterization is not None:
-            properties["pruning_target_parameterization"] = arm_to_dict(
-                oc.pruning_target_parameterization
-            )
-        if Keys.LLM_MESSAGES in properties:
-            properties[Keys.LLM_MESSAGES] = [
-                dataclasses.asdict(m) if isinstance(m, LLMMessage) else m
-                for m in properties[Keys.LLM_MESSAGES]
-            ]
+        properties = prepare_experiment_properties_for_storage(experiment)
 
         # pyre-ignore[9]: Expected `Base` for 1st...yping.Type[Experiment]`.
         experiment_class: type[SQAExperiment] = self.config.class_to_sqa_class[

ax/storage/sqa_store/save.py

@@ -29,7 +29,10 @@
 from ax.storage.sqa_store import validation as _validation_listeners  # noqa: F401
 from ax.storage.sqa_store.db import session_scope, SQABase
 from ax.storage.sqa_store.decoder import Decoder
-from ax.storage.sqa_store.encoder import Encoder
+from ax.storage.sqa_store.encoder import (
+    Encoder,
+    prepare_experiment_properties_for_storage,
+)
 from ax.storage.sqa_store.sqa_classes import (
     SQAData,
     SQAExperiment,
@@ -540,10 +543,14 @@ def update_properties_on_experiment(
 
     exp_id = _assert_experiment_saved(experiment_with_updated_properties)
 
+    properties = prepare_experiment_properties_for_storage(
+        experiment_with_updated_properties
+    )
+
     with session_scope() as session:
         session.query(exp_sqa_class).filter_by(id=exp_id).update(
             {
-                "properties": experiment_with_updated_properties._properties,
+                "properties": properties,
             }
         )
 

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -34,6 +34,7 @@
 from ax.core.experiment import Experiment
 from ax.core.experiment_status import ExperimentStatus
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage
 from ax.core.metric import Metric
 from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
@@ -2863,6 +2864,29 @@ def test_set_immutable_search_space_and_opt_config(self) -> None:
         loaded_experiment = load_experiment(experiment.name)
         self.assertTrue(loaded_experiment.immutable_search_space_and_opt_config)
 
+    def test_update_properties_on_experiment_with_llm_messages(self) -> None:
+        """Test that LLMMessage objects in experiment properties are correctly
+        serialized through the incremental update_properties_on_experiment path,
+        not just the full experiment_to_sqa path."""
+        experiment = get_experiment_with_batch_trial()
+        save_experiment(experiment)
+
+        messages = [
+            LLMMessage(role="system", content="You are helpful."),
+            LLMMessage(role="user", content="Hello", metadata={"key": "val"}),
+        ]
+        experiment.llm_messages = messages
+        update_properties_on_experiment(
+            experiment_with_updated_properties=experiment,
+        )
+
+        loaded_experiment = load_experiment(experiment.name)
+        self.assertEqual(len(loaded_experiment.llm_messages), 2)
+        self.assertEqual(loaded_experiment.llm_messages[0].role, "system")
+        self.assertEqual(loaded_experiment.llm_messages[0].content, "You are helpful.")
+        self.assertEqual(loaded_experiment.llm_messages[1].role, "user")
+        self.assertEqual(loaded_experiment.llm_messages[1].metadata, {"key": "val"})
+
     def test_update_properties_on_trial(self) -> None:
         experiment = get_experiment_with_batch_trial()
         self.assertNotIn("foo", experiment.trials[0]._properties)