Enable json serialization of ScalarizedOutcomeConstraint (#5025)

Summary:

Adds JSON serialization/deserialization support for ScalarizedOutcomeConstraint in Ax's JSON storage.

= Background

ScalarizedOutcomeConstraint differs from OutcomeConstraint in that it holds a list of metrics and weights instead of a single metric. Accessing .metric on a ScalarizedOutcomeConstraint raises NotImplementedError, so the existing outcome_constraint_to_dict encoder cannot be reused.

= Solution

This diff implements a dedicated encoder (scalarized_outcome_constraint_to_dict) following the same pattern as ScalarizedObjective.

= Motivation

This issue was discovered when serializing an experiment loaded from the Ax DB via object_to_json. The experiment's historical GeneratorRun objects contained ScalarizedOutcomeConstraint instances (set by Bayesian optimization), which caused serialization to fail.

**Note:** The SQA (database) store already handles this type; the JSON store was the only gap.

Reviewed By: bletham

Differential Revision: D96486632

Author: metactc

ax/storage/json_store/encoders.py

@@ -26,7 +26,7 @@
     OptimizationConfig,
     PreferenceOptimizationConfig,
 )
-from ax.core.outcome_constraint import OutcomeConstraint
+from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint
 from ax.core.parameter import (
     ChoiceParameter,
     DerivedParameter,
@@ -362,6 +362,21 @@ def outcome_constraint_to_dict(outcome_constraint: OutcomeConstraint) -> dict[st
     }
 
 
+def scalarized_outcome_constraint_to_dict(
+    outcome_constraint: ScalarizedOutcomeConstraint,
+) -> dict[str, Any]:
+    """Convert Ax scalarized outcome constraint to a dictionary."""
+    metrics = [Metric(name=name) for name, _ in outcome_constraint.metric_weights]
+    return {
+        "__type": outcome_constraint.__class__.__name__,
+        "metrics": metrics,
+        "weights": outcome_constraint.weights,
+        "op": outcome_constraint.op,
+        "bound": outcome_constraint.bound,
+        "relative": outcome_constraint.relative,
+    }
+
+
 def optimization_config_to_dict(
     optimization_config: OptimizationConfig,
 ) -> dict[str, Any]:

ax/storage/json_store/registry.py

@@ -47,7 +47,11 @@
     OptimizationConfig,
     PreferenceOptimizationConfig,
 )
-from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint
+from ax.core.outcome_constraint import (
+    ObjectiveThreshold,
+    OutcomeConstraint,
+    ScalarizedOutcomeConstraint,
+)
 from ax.core.parameter import (
     ChoiceParameter,
     DerivedParameter,
@@ -159,6 +163,7 @@
     range_parameter_to_dict,
     runner_to_dict,
     scalarized_objective_to_dict,
+    scalarized_outcome_constraint_to_dict,
     search_space_to_dict,
     surrogate_to_dict,
     threshold_early_stopping_strategy_to_dict,
@@ -255,6 +260,7 @@
     OptimizationConfig: optimization_config_to_dict,
     OrEarlyStoppingStrategy: logical_early_stopping_strategy_to_dict,
     OutcomeConstraint: outcome_constraint_to_dict,
+    ScalarizedOutcomeConstraint: scalarized_outcome_constraint_to_dict,
     ParameterConstraint: parameter_constraint_to_dict,
     pathlib.Path: pathlib_to_dict,
     pathlib.PurePath: pathlib_to_dict,
@@ -410,6 +416,7 @@
     "ReductionCriterion": ReductionCriterion,
     "Round": Round,
     "ScalarizedObjective": ScalarizedObjective,
+    "ScalarizedOutcomeConstraint": ScalarizedOutcomeConstraint,
     "SchedulerOptions": OrchestratorOptions,  # DEPRECATED; backward compatibility
     "SearchSpace": SearchSpace,
     "SimTrial": SimTrial,

ax/storage/json_store/tests/test_json_store.py

@@ -153,6 +153,7 @@
     get_percentile_early_stopping_strategy_with_non_objective_metric_signature,
     get_range_parameter,
     get_scalarized_objective,
+    get_scalarized_outcome_constraint,
     get_search_space,
     get_sorted_choice_parameter,
     get_sum_constraint1,
@@ -391,6 +392,7 @@
     ("ParameterConstraint", get_parameter_constraint),
     ("RangeParameter", get_range_parameter),
     ("ScalarizedObjective", get_scalarized_objective),
+    ("ScalarizedOutcomeConstraint", get_scalarized_outcome_constraint),
     ("OrchestratorOptions", get_default_orchestrator_options),
     ("OrchestratorOptions", get_orchestrator_options_batch_trial),
     ("SearchSpace", get_search_space),