Remove special logic for BenchmarkProblems in JSON serialization (#…

…2611)

Summary:
Pull Request resolved: #2611

These recently became dataclasses. Dataclasses are generically handled by Ax's JSON encoders and decoders, so the special logic for serializing BenchmarkProblems can be removed.

Reviewed By: saitcakmak

Differential Revision: D60353847

fbshipit-source-id: 62eac8a1273da67bf7def3ee5945ce16c57f9dcf

ax/storage/json_store/encoders.py

@@ -11,10 +11,6 @@
 from pathlib import Path
 from typing import Any, Dict, Type
 
-from ax.benchmark.benchmark_problem import (
-    BenchmarkProblem,
-    MultiObjectiveBenchmarkProblem,
-)
 from ax.benchmark.problems.hpo.torchvision import PyTorchCNNTorchvisionBenchmarkProblem
 from ax.core import ObservationFeatures
 from ax.core.arm import Arm
@@ -133,32 +129,6 @@ def batch_to_dict(batch: BatchTrial) -> Dict[str, Any]:
     }
 
 
-def benchmark_problem_to_dict(benchmark_problem: BenchmarkProblem) -> Dict[str, Any]:
-    """Convert Ax benchmark problem to a dictionary."""
-    return {
-        "__type": benchmark_problem.__class__.__name__,
-        "name": benchmark_problem.name,
-        "search_space": benchmark_problem.search_space,
-        "optimization_config": benchmark_problem.optimization_config,
-        "runner": benchmark_problem.runner,
-        "num_trials": benchmark_problem.num_trials,
-        "is_noiseless": benchmark_problem.is_noiseless,
-        "observe_noise_stds": benchmark_problem.observe_noise_stds,
-        "has_ground_truth": benchmark_problem.has_ground_truth,
-        "tracking_metrics": benchmark_problem.tracking_metrics,
-        "optimal_value": benchmark_problem.optimal_value,
-    }
-
-
-def multi_objective_benchmark_problem_to_dict(
-    moo_benchmark_problem: MultiObjectiveBenchmarkProblem,
-) -> Dict[str, Any]:
-    """Convert Ax multi-objective benchmark problem to a dictionary."""
-    result = benchmark_problem_to_dict(moo_benchmark_problem)
-    result["reference_point"] = moo_benchmark_problem.reference_point
-    return result
-
-
 def trial_to_dict(trial: Trial) -> Dict[str, Any]:
     """Convert Ax trial to a dictionary."""
     return {

ax/storage/json_store/registry.py

@@ -116,7 +116,6 @@
 from ax.storage.json_store.encoders import (
     arm_to_dict,
     batch_to_dict,
-    benchmark_problem_to_dict,
     best_model_selector_to_dict,
     botorch_component_to_dict,
     botorch_model_to_dict,
@@ -136,7 +135,6 @@
     map_key_info_to_dict,
     metric_to_dict,
     model_spec_to_dict,
-    multi_objective_benchmark_problem_to_dict,
     multi_objective_optimization_config_to_dict,
     multi_objective_to_dict,
     multi_type_experiment_to_dict,
@@ -186,7 +184,6 @@
     AugmentedHartmann6Metric: metric_to_dict,
     BatchTrial: batch_to_dict,
     BenchmarkMetric: metric_to_dict,
-    BenchmarkProblem: benchmark_problem_to_dict,
     BoTorchModel: botorch_model_to_dict,
     BotorchTestProblemRunner: runner_to_dict,
     BraninMetric: metric_to_dict,
@@ -220,7 +217,6 @@
     MinimumPreferenceOccurances: transition_criterion_to_dict,
     ModelSpec: model_spec_to_dict,
     MultiObjective: multi_objective_to_dict,
-    MultiObjectiveBenchmarkProblem: multi_objective_benchmark_problem_to_dict,
     MultiObjectiveOptimizationConfig: multi_objective_optimization_config_to_dict,
     MultiTypeExperiment: multi_type_experiment_to_dict,
     Normalize: botorch_component_to_dict,