Skip to content

Commit 18fa6ec

Browse files
metactcfacebook-github-bot
authored andcommitted
Enable json serialization of ScalarizedOutcomeConstraint
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
1 parent e2056d2 commit 18fa6ec

File tree

3 files changed

+25
-2
lines changed

3 files changed

+25
-2
lines changed

ax/storage/json_store/encoders.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
OptimizationConfig,
2626
PreferenceOptimizationConfig,
2727
)
28-
from ax.core.outcome_constraint import OutcomeConstraint
28+
from ax.core.outcome_constraint import OutcomeConstraint, ScalarizedOutcomeConstraint
2929
from ax.core.parameter import (
3030
ChoiceParameter,
3131
DerivedParameter,
@@ -300,6 +300,20 @@ def outcome_constraint_to_dict(outcome_constraint: OutcomeConstraint) -> dict[st
300300
}
301301

302302

303+
def scalarized_outcome_constraint_to_dict(
304+
outcome_constraint: ScalarizedOutcomeConstraint,
305+
) -> dict[str, Any]:
306+
"""Convert Ax scalarized outcome constraint to a dictionary."""
307+
return {
308+
"__type": outcome_constraint.__class__.__name__,
309+
"metrics": outcome_constraint.metrics,
310+
"weights": outcome_constraint.weights,
311+
"op": outcome_constraint.op,
312+
"bound": outcome_constraint.bound,
313+
"relative": outcome_constraint.relative,
314+
}
315+
316+
303317
def optimization_config_to_dict(
304318
optimization_config: OptimizationConfig,
305319
) -> dict[str, Any]:

ax/storage/json_store/registry.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,11 @@
4242
OptimizationConfig,
4343
PreferenceOptimizationConfig,
4444
)
45-
from ax.core.outcome_constraint import ObjectiveThreshold, OutcomeConstraint
45+
from ax.core.outcome_constraint import (
46+
ObjectiveThreshold,
47+
OutcomeConstraint,
48+
ScalarizedOutcomeConstraint,
49+
)
4650
from ax.core.parameter import (
4751
ChoiceParameter,
4852
DerivedParameter,
@@ -152,6 +156,7 @@
152156
range_parameter_to_dict,
153157
runner_to_dict,
154158
scalarized_objective_to_dict,
159+
scalarized_outcome_constraint_to_dict,
155160
search_space_to_dict,
156161
surrogate_to_dict,
157162
threshold_early_stopping_strategy_to_dict,
@@ -241,6 +246,7 @@
241246
OptimizationConfig: optimization_config_to_dict,
242247
OrEarlyStoppingStrategy: logical_early_stopping_strategy_to_dict,
243248
OutcomeConstraint: outcome_constraint_to_dict,
249+
ScalarizedOutcomeConstraint: scalarized_outcome_constraint_to_dict,
244250
ParameterConstraint: parameter_constraint_to_dict,
245251
pathlib.Path: pathlib_to_dict,
246252
pathlib.PurePath: pathlib_to_dict,
@@ -389,6 +395,7 @@
389395
"ReductionCriterion": ReductionCriterion,
390396
"Round": Round,
391397
"ScalarizedObjective": ScalarizedObjective,
398+
"ScalarizedOutcomeConstraint": ScalarizedOutcomeConstraint,
392399
"SchedulerOptions": OrchestratorOptions, # DEPRECATED; backward compatibility
393400
"SearchSpace": SearchSpace,
394401
"SimTrial": SimTrial,

ax/storage/json_store/tests/test_json_store.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@
148148
get_percentile_early_stopping_strategy_with_non_objective_metric_signature,
149149
get_range_parameter,
150150
get_scalarized_objective,
151+
get_scalarized_outcome_constraint,
151152
get_search_space,
152153
get_sorted_choice_parameter,
153154
get_sum_constraint1,
@@ -386,6 +387,7 @@
386387
("ParameterConstraint", get_parameter_constraint),
387388
("RangeParameter", get_range_parameter),
388389
("ScalarizedObjective", get_scalarized_objective),
390+
("ScalarizedOutcomeConstraint", get_scalarized_outcome_constraint),
389391
("OrchestratorOptions", get_default_orchestrator_options),
390392
("OrchestratorOptions", get_orchestrator_options_batch_trial),
391393
("SearchSpace", get_search_space),

0 commit comments

Comments
 (0)