Add storage support for RangeParameter.step_size (JSON + SQA)

Summary:
Teaches both storage backends to persist and restore the new `step_size` field on `RangeParameter`. Second diff in the step_size unification stack (builds on the native `step_size` support added in D107274057).

The underlying SQA column (`SQAParameter.step_size`) and the corresponding EntAE schema were added and landed separately (D107108021, D107112306), so this diff is pure encoder/decoder logic — it does not touch the schema.

SQA: the encoder writes `parameter.step_size` to the column (and continues writing the legacy `digits` column for now; `digits` is dropped in a later diff). The decoder prefers `step_size` and falls back to the legacy `digits` column for rows written by older code, passing only one to the constructor (which rejects both being set; it converts `digits` to `step_size` internally).

JSON: `range_parameter_to_dict` emits `step_size` alongside `digits`. There is no custom JSON decoder for `RangeParameter` — the registry maps it directly to the class, and `ax_class_from_json_dict` splats the serialized dict straight into the constructor. So `step_size` flows through automatically, and legacy blobs carrying only `digits` still decode via the constructor's back-compat handling. No decoder change is required for JSON.

Differential Revision: D107282941

Author: saitcakmak

ax/storage/json_store/encoders.py

@@ -201,6 +201,7 @@ def range_parameter_to_dict(parameter: RangeParameter) -> dict[str, Any]:
         "log_scale": parameter.log_scale,
         "logit_scale": parameter.logit_scale,
         "digits": parameter.digits,
+        "step_size": parameter.step_size,
         "is_fidelity": parameter.is_fidelity,
         "target_value": parameter.target_value,
     }

ax/storage/json_store/tests/test_json_store.py

@@ -55,7 +55,7 @@
     OptimizationConfig,
     PreferenceOptimizationConfig,
 )
-from ax.core.parameter import ChoiceParameter, ParameterType
+from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.runner import Runner
 from ax.exceptions.core import AxStorageWarning, UnsupportedError
@@ -398,6 +398,17 @@
     ("ParameterConstraint", get_parameter_constraint),
     ("ParameterConstraint", get_equality_parameter_constraint),
     ("RangeParameter", get_range_parameter),
+    (
+        "RangeParameter",
+        partial(
+            RangeParameter,
+            name="x",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=1.0,
+            step_size=0.1,
+        ),
+    ),
     ("ScalarizedObjective", get_scalarized_objective),
     ("ScalarizedOutcomeConstraint", get_scalarized_outcome_constraint),
     ("OrchestratorOptions", get_default_orchestrator_options),
@@ -1942,6 +1953,32 @@ def test_multi_objective_from_json_warning(self) -> None:
             any("Found unexpected kwargs" in warning for warning in cm.output)
         )
 
+    def test_range_parameter_legacy_digits_blob_decodes(self) -> None:
+        # A legacy blob has "digits" but no "step_size" key. It must still
+        # decode (the constructor accepts digits for back-compat).
+        legacy_blob = {
+            "__type": "RangeParameter",
+            "name": "x",
+            "parameter_type": {"__type": "ParameterType", "name": "FLOAT"},
+            "lower": 0.0,
+            "upper": 1.0,
+            "log_scale": False,
+            "logit_scale": False,
+            "digits": 2,
+            "is_fidelity": False,
+            "target_value": None,
+        }
+        decoded = object_from_json(
+            legacy_blob,
+            decoder_registry=CORE_DECODER_REGISTRY,
+            class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+        )
+        self.assertIsInstance(decoded, RangeParameter)
+        self.assertEqual(decoded.digits, 2)
+        self.assertIsNone(decoded.step_size)
+        # Rounding behavior from digits=2 is preserved.
+        self.assertEqual(decoded.cast(0.123), 0.12)
+
     def test_choice_parameter_bypass_cardinality_check_encode_failure(self) -> None:
         choice_parameter = ChoiceParameter(
             name="test_choice",

ax/storage/sqa_store/decoder.py

@@ -470,13 +470,21 @@ def parameter_from_sqa(self, parameter_sqa: SQAParameter) -> Parameter:
                     "`dependents` unexpectedly non-null on range parameter "
                     f"{parameter_sqa.name}."
                 )
+            # Prefer the newer ``step_size`` column; fall back to the legacy
+            # ``digits`` column for rows written by older code. Only one may be
+            # passed to the constructor (it rejects both being set). The
+            # constructor converts ``digits`` to ``step_size`` internally.
+            # We should never have the two together in the DB, so this is extra.
+            step_size = parameter_sqa.step_size
+            digits = None if step_size is not None else parameter_sqa.digits
             parameter = RangeParameter(
                 name=parameter_sqa.name,
                 parameter_type=parameter_sqa.parameter_type,
                 lower=float(none_throws(parameter_sqa.lower)),
                 upper=float(none_throws(parameter_sqa.upper)),
                 log_scale=parameter_sqa.log_scale or False,
-                digits=parameter_sqa.digits,
+                digits=digits,
+                step_size=float(step_size) if step_size is not None else None,
                 is_fidelity=parameter_sqa.is_fidelity or False,
                 target_value=parameter_sqa.target_value,
                 backfill_value=parameter_sqa.backfill_value,

ax/storage/sqa_store/encoder.py

@@ -308,6 +308,7 @@ def parameter_to_sqa(self, parameter: Parameter) -> SQAParameter:
                 upper=float(parameter.upper),
                 log_scale=parameter.log_scale,
                 digits=parameter.digits,
+                step_size=parameter.step_size,
                 is_fidelity=parameter.is_fidelity,
                 target_value=parameter.target_value,
                 dependents=parameter.dependents if parameter.is_hierarchical else None,

ax/storage/sqa_store/tests/test_sqa_store.py

@@ -1880,6 +1880,41 @@ def test_logit_scale(self) -> None:
                 )
             )
 
+    def test_step_size_round_trip(self) -> None:
+        parameter = RangeParameter(
+            name="foo",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=1.0,
+            step_size=0.1,
+        )
+        sqa_parameter = self.encoder.parameter_to_sqa(parameter)
+        self.assertEqual(sqa_parameter.step_size, 0.1)
+        self.assertIsNone(sqa_parameter.digits)
+        decoded = self.decoder.parameter_from_sqa(sqa_parameter)
+        self.assertEqual(decoded, parameter)
+        self.assertEqual(assert_is_instance(decoded, RangeParameter).step_size, 0.1)
+
+    def test_legacy_digits_row_decodes_via_fallback(self) -> None:
+        # A legacy row carries ``digits`` but no ``step_size``. The decoder
+        # falls back to ``digits``, which the constructor converts internally.
+        parameter = RangeParameter(
+            name="foo",
+            parameter_type=ParameterType.FLOAT,
+            lower=0.0,
+            upper=1.0,
+            digits=2,
+        )
+        sqa_parameter = self.encoder.parameter_to_sqa(parameter)
+        # Simulate an old row: clear step_size, keep digits populated.
+        sqa_parameter.step_size = None
+        sqa_parameter.digits = 2
+        decoded = assert_is_instance(
+            self.decoder.parameter_from_sqa(sqa_parameter), RangeParameter
+        )
+        # The decoded parameter round-trips its rounding behavior.
+        self.assertEqual(decoded.cast(0.123), parameter.cast(0.123))
+
     def test_bypass_cardinality_check(self) -> None:
         choice_parameter = ChoiceParameter(
             name="test_choice",