cast arms' parameter values when decoding from json/sqa (facebook#4853)

Summary:
Pull Request resolved: facebook#4853

See title. This casts arms' parameter values to the proper type when decoding from json/sqa. This fixes an issue where SQ specified via the UI can result in Float parameters being loaded as ints from SQA.

This led to issues with imported QEs, where the imported SQ parameter had a value of 1.0, but the signature of the new SQ did not match the signature of the existing SQ value of 1.

Reviewed By: Cesar-Cardoso

Differential Revision: D92182397

fbshipit-source-id: 80bc928d43d93306acce268d23626fcec66096ff

Author: sdaulton

ax/storage/json_store/decoder.py

@@ -22,6 +22,7 @@
 import pandas as pd
 import torch
 from ax.adapter.registry import GeneratorRegistryBase
+from ax.core.arm import Arm
 from ax.core.base_trial import BaseTrial
 from ax.core.data import Data
 from ax.core.experiment import Experiment
@@ -47,6 +48,7 @@
 from ax.generators.torch.botorch_modular.surrogate import Surrogate, SurrogateSpec
 from ax.generators.torch.botorch_modular.utils import ModelConfig
 from ax.storage.json_store.decoders import (
+    _cast_parameter_value,
     batch_trial_from_json,
     botorch_component_from_json,
     tensor_from_json,
@@ -72,6 +74,24 @@
 logger: Logger = get_logger(__name__)
 
 
+def _cast_arm_parameters(arm: Arm, search_space: SearchSpace) -> None:
+    """Cast arm parameter values to the appropriate Python type.
+
+    This is necessary because JSON may deserialize values as different types
+    (e.g., ints as floats). This function modifies the arm in place.
+
+    Args:
+        arm: The arm whose parameter values should be cast.
+        search_space: The search space containing parameter type information.
+    """
+    for param_name, param_value in arm._parameters.items():
+        if param_name in search_space.parameters:
+            parameter = search_space.parameters[param_name]
+            arm._parameters[param_name] = _cast_parameter_value(
+                param_value, parameter.parameter_type
+            )
+
+
 def _raise_on_legacy_callable_refs(kwarg_dict: dict[str, Any]) -> dict[str, Any]:
     """Returns kwarg_dict unchanged if no legacy callable refs are present.
 
@@ -759,11 +779,17 @@ def _load_experiment_info(
     )
     for trial in exp._trials.values():
         for arm in trial.arms:
+            # Cast arm parameter values to the appropriate type based on the
+            # search space parameter types. This is necessary because JSON may
+            # deserialize values as different types (e.g., ints as floats).
+            _cast_arm_parameters(arm, exp.search_space)
             exp._register_arm(arm)
         if trial.ttl_seconds is not None:
             exp._trials_have_ttl = True
     if exp.status_quo is not None:
         sq = none_throws(exp.status_quo)
+        # Cast status_quo arm parameter values as well.
+        _cast_arm_parameters(sq, exp.search_space)
         exp._register_arm(sq)
 
 

ax/storage/json_store/decoders.py

@@ -26,6 +26,7 @@
 from ax.core.parameter import (
     ChoiceParameter,
     FixedParameter,
+    PARAMETER_PYTHON_TYPE_MAP,
     ParameterType,
     TParamValue,
 )
@@ -93,6 +94,27 @@ def string_to_parameter_value(s: str, parameter_type: ParameterType) -> TParamVa
         return s
 
 
+def _cast_parameter_value(
+    value: TParamValue, parameter_type: ParameterType
+) -> TParamValue:
+    """Cast a parameter value to the appropriate Python type based on parameter_type.
+
+    This is necessary because JSON may deserialize values as different types
+    (e.g., ints as floats).
+
+    Args:
+        value: The value to cast.
+        parameter_type: The ParameterType to cast to.
+
+    Returns:
+        The value cast to the appropriate Python type.
+    """
+    if value is None:
+        return None
+    python_type = PARAMETER_PYTHON_TYPE_MAP[parameter_type]
+    return python_type(value)
+
+
 def batch_trial_from_json(
     experiment: core.experiment.Experiment,
     index: int,

ax/storage/json_store/tests/test_json_store.py

@@ -1725,6 +1725,145 @@ def test_choice_parameter_backward_compatibility_sort_values(self) -> None:
             )
         )
 
+    def test_arm_parameter_values_cast_to_parameter_type(self) -> None:
+        """Test that arm parameter values are cast to the appropriate type on load."""
+        from ax.core.arm import Arm
+        from ax.core.experiment import Experiment
+        from ax.core.parameter import RangeParameter
+        from ax.core.search_space import SearchSpace
+        from ax.storage.json_store.encoder import object_to_json
+
+        # Create an experiment with INT parameters
+        search_space = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="x",
+                    parameter_type=ParameterType.INT,
+                    lower=0,
+                    upper=10,
+                ),
+                RangeParameter(
+                    name="y",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=1.0,
+                ),
+            ]
+        )
+
+        experiment = Experiment(
+            name="test_experiment",
+            search_space=search_space,
+            status_quo=Arm(parameters={"x": 5, "y": 0.5}, name="status_quo"),
+        )
+
+        # Add a trial with an arm
+        trial = experiment.new_trial()
+        trial.add_arm(Arm(parameters={"x": 3, "y": 0.3}))
+
+        # Encode the experiment to JSON
+        experiment_json = object_to_json(
+            experiment,
+            encoder_registry=CORE_ENCODER_REGISTRY,
+            class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+        )
+
+        # Manually modify the JSON to simulate float values for INT parameters
+        # (as could happen when loading from external sources)
+        for arm_json in experiment_json["trials"][0]["generator_run"]["arms"]:
+            arm_json["parameters"]["x"] = 3.0  # float instead of int
+        experiment_json["status_quo"]["parameters"]["x"] = 5.0  # float instead of int
+
+        # Decode the experiment from JSON
+        loaded_experiment = object_from_json(
+            experiment_json,
+            decoder_registry=CORE_DECODER_REGISTRY,
+            class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+        )
+
+        # Check that arm parameter values are cast to the correct type
+        loaded_arm = list(loaded_experiment.trials[0].arms)[0]
+        self.assertEqual(loaded_arm.parameters["x"], 3)
+        self.assertIs(type(loaded_arm.parameters["x"]), int)
+        self.assertEqual(loaded_arm.parameters["y"], 0.3)
+        self.assertIs(type(loaded_arm.parameters["y"]), float)
+
+        # Check that status_quo parameter values are cast to the correct type
+        status_quo = loaded_experiment.status_quo
+        self.assertIsNotNone(status_quo)
+        self.assertEqual(status_quo.parameters["x"], 5)
+        self.assertIs(type(status_quo.parameters["x"]), int)
+        self.assertEqual(status_quo.parameters["y"], 0.5)
+        self.assertIs(type(status_quo.parameters["y"]), float)
+
+    def test_cast_parameter_value_all_types(self) -> None:
+        """Test _cast_parameter_value handles all parameter types correctly."""
+        from ax.storage.json_store.decoders import _cast_parameter_value
+
+        # Test INT casting
+        self.assertEqual(_cast_parameter_value(3.0, ParameterType.INT), 3)
+        self.assertIs(type(_cast_parameter_value(3.0, ParameterType.INT)), int)
+        self.assertEqual(_cast_parameter_value(3, ParameterType.INT), 3)
+        self.assertIs(type(_cast_parameter_value(3, ParameterType.INT)), int)
+
+        # Test FLOAT casting
+        self.assertEqual(_cast_parameter_value(3, ParameterType.FLOAT), 3.0)
+        self.assertIs(type(_cast_parameter_value(3, ParameterType.FLOAT)), float)
+        self.assertEqual(_cast_parameter_value(3.5, ParameterType.FLOAT), 3.5)
+        self.assertIs(type(_cast_parameter_value(3.5, ParameterType.FLOAT)), float)
+
+        # Test BOOL casting
+        self.assertEqual(_cast_parameter_value(1, ParameterType.BOOL), True)
+        self.assertIs(type(_cast_parameter_value(1, ParameterType.BOOL)), bool)
+        self.assertEqual(_cast_parameter_value(0, ParameterType.BOOL), False)
+        self.assertIs(type(_cast_parameter_value(0, ParameterType.BOOL)), bool)
+        self.assertEqual(_cast_parameter_value(True, ParameterType.BOOL), True)
+        self.assertIs(type(_cast_parameter_value(True, ParameterType.BOOL)), bool)
+
+        # Test STRING casting
+        self.assertEqual(_cast_parameter_value("test", ParameterType.STRING), "test")
+        self.assertIs(type(_cast_parameter_value("test", ParameterType.STRING)), str)
+        self.assertEqual(_cast_parameter_value(123, ParameterType.STRING), "123")
+        self.assertIs(type(_cast_parameter_value(123, ParameterType.STRING)), str)
+
+        # Test None handling
+        self.assertIsNone(_cast_parameter_value(None, ParameterType.INT))
+        self.assertIsNone(_cast_parameter_value(None, ParameterType.FLOAT))
+        self.assertIsNone(_cast_parameter_value(None, ParameterType.BOOL))
+        self.assertIsNone(_cast_parameter_value(None, ParameterType.STRING))
+
+    def test_cast_arm_parameters_skips_unknown_params(self) -> None:
+        """Test that _cast_arm_parameters skips parameters not in search space."""
+        from ax.core.arm import Arm
+        from ax.core.parameter import RangeParameter
+        from ax.core.search_space import SearchSpace
+        from ax.storage.json_store.decoder import _cast_arm_parameters
+
+        search_space = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="x",
+                    parameter_type=ParameterType.INT,
+                    lower=0,
+                    upper=10,
+                ),
+            ]
+        )
+
+        # Create an arm with a parameter that's not in the search space
+        arm = Arm(parameters={"x": 3.0, "unknown_param": "some_value"})
+
+        # Cast should work without error and only cast known parameters
+        _cast_arm_parameters(arm, search_space)
+
+        # x should be cast to int
+        self.assertEqual(arm.parameters["x"], 3)
+        self.assertIs(type(arm.parameters["x"]), int)
+
+        # unknown_param should remain unchanged
+        self.assertEqual(arm.parameters["unknown_param"], "some_value")
+        self.assertIs(type(arm.parameters["unknown_param"]), str)
+
     def test_surrogate_spec_backwards_compatibility(self) -> None:
         # This is an invalid example that has both deprecated args
         # and model config specified. Deprecated args will be ignored.

ax/storage/sqa_store/decoder.py

@@ -57,6 +57,7 @@
 from ax.exceptions.storage import JSONDecodeError, SQADecodeError
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.storage.json_store.decoder import _DEPRECATED_GENERATOR_KWARGS, object_from_json
+from ax.storage.json_store.decoders import _cast_parameter_value
 from ax.storage.sqa_store.db import session_scope
 from ax.storage.sqa_store.sqa_classes import (
     SQAAbandonedArm,
@@ -90,6 +91,24 @@
 logger: Logger = get_logger(__name__)
 
 
+def _cast_arm_parameters(arm: Arm, search_space: SearchSpace) -> None:
+    """Cast arm parameter values to the appropriate Python type.
+
+    This is necessary because SQA may deserialize values as different types
+    (e.g., ints as floats). This function modifies the arm in place.
+
+    Args:
+        arm: The arm whose parameter values should be cast.
+        search_space: The search space containing parameter type information.
+    """
+    for param_name, param_value in arm._parameters.items():
+        if param_name in search_space.parameters:
+            parameter = search_space.parameters[param_name]
+            arm._parameters[param_name] = _cast_parameter_value(
+                param_value, parameter.parameter_type
+            )
+
+
 class Decoder:
     """Class that contains methods for loading an Ax experiment from SQLAlchemy.
 
@@ -378,9 +397,15 @@ def experiment_from_sqa(
             if trial.ttl_seconds is not None:
                 experiment._trials_have_ttl = True
             for arm in trial.arms:
+                # Cast arm parameter values to the appropriate type based on the
+                # search space parameter types. This is necessary because SQA may
+                # deserialize values as different types (e.g., ints as floats).
+                _cast_arm_parameters(arm, experiment.search_space)
                 experiment._register_arm(arm)
         if experiment.status_quo is not None:
             sq = none_throws(experiment.status_quo)
+            # Cast status_quo arm parameter values as well.
+            _cast_arm_parameters(sq, experiment.search_space)
             experiment._register_arm(sq)
         experiment._time_created = experiment_sqa.time_created
         experiment._experiment_type = self.get_enum_name(

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -43,6 +43,7 @@
 from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.runner import Runner
+from ax.core.search_space import SearchSpace
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
 from ax.core.types import ComparisonOp
@@ -350,6 +351,77 @@ def test_experiment_save_load(self) -> None:
             loaded_experiment = load_experiment(exp.name)
             self.assertEqual(loaded_experiment, exp)
 
+    def test_arm_parameter_values_cast_to_parameter_type(self) -> None:
+        """Test that arm parameter values are cast to the appropriate type on load.
+
+        This is important because SQA may deserialize values as different types
+        (e.g., ints as floats). This test directly modifies the SQA objects to
+        simulate float values for INT parameters, then verifies that the decoder
+        casts them back to the correct types.
+        """
+        # Create an experiment with INT and FLOAT parameters
+        search_space = SearchSpace(
+            parameters=[
+                RangeParameter(
+                    name="x",
+                    parameter_type=ParameterType.INT,
+                    lower=0,
+                    upper=10,
+                ),
+                RangeParameter(
+                    name="y",
+                    parameter_type=ParameterType.FLOAT,
+                    lower=0.0,
+                    upper=1.0,
+                ),
+            ]
+        )
+
+        experiment = Experiment(
+            name="test_arm_param_cast_experiment_sqa",
+            search_space=search_space,
+            status_quo=Arm(parameters={"x": 5, "y": 0.5}, name="status_quo"),
+        )
+
+        # Add a trial with an arm
+        trial = experiment.new_trial()
+        trial.add_arm(Arm(parameters={"x": 3, "y": 0.3}))
+
+        # Encode the experiment to SQA
+        experiment_sqa = self.encoder.experiment_to_sqa(experiment)
+
+        # Manually modify the SQA objects to simulate float values for INT parameters
+        # This simulates what can happen when loading from external sources or
+        # when JSON deserialization converts ints to floats.
+        # Modify status_quo parameters
+        none_throws(experiment_sqa.status_quo_parameters)["x"] = (
+            5.0  # float instead of int
+        )
+
+        # Modify arm parameters in the trial's generator run
+        for trial_sqa in experiment_sqa.trials:
+            for gr_sqa in trial_sqa.generator_runs:
+                for arm_sqa in gr_sqa.arms:
+                    arm_sqa.parameters["x"] = 3.0  # float instead of int
+
+        # Decode the experiment from SQA
+        loaded_experiment = self.decoder.experiment_from_sqa(experiment_sqa)
+
+        # Check that arm parameter values are cast to the correct type
+        loaded_arm = list(loaded_experiment.trials[0].arms)[0]
+        self.assertEqual(loaded_arm.parameters["x"], 3)
+        self.assertIs(type(loaded_arm.parameters["x"]), int)
+        self.assertEqual(loaded_arm.parameters["y"], 0.3)
+        self.assertIs(type(loaded_arm.parameters["y"]), float)
+
+        # Check that status_quo parameter values are cast to the correct type
+        status_quo = loaded_experiment.status_quo
+        self.assertIsNotNone(status_quo)
+        self.assertEqual(none_throws(status_quo).parameters["x"], 5)
+        self.assertIs(type(none_throws(status_quo).parameters["x"]), int)
+        self.assertEqual(none_throws(status_quo).parameters["y"], 0.5)
+        self.assertIs(type(none_throws(status_quo).parameters["y"]), float)
+
     def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
         aux_experiment = Experiment(
             name="test_aux_exp_in_SQAStoreTest",