Source code for ax.storage.json_store.decoder

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

import datetime
import enum
from collections import OrderedDict
from typing import Any, Dict, List

import pandas as pd
from ax.core.base_trial import BaseTrial
from ax.core.data import Data  # noqa F401
from ax.core.experiment import Experiment
from ax.core.generator_run import GeneratorRun
from ax.core.parameter import Parameter
from ax.core.parameter_constraint import (
    OrderConstraint,
    ParameterConstraint,
    SumConstraint,
)
from ax.core.search_space import SearchSpace
from ax.core.simple_experiment import (
    SimpleExperiment,
    unimplemented_evaluation_function,
)
from ax.exceptions.storage import JSONDecodeError
from ax.storage.json_store.decoders import batch_trial_from_json, trial_from_json
from ax.storage.json_store.registry import DECODER_REGISTRY


[docs]def object_from_json(object_json: Any) -> Any: """Recursively load objects from a JSON-serializable dictionary.""" if type(object_json) in (str, int, float, bool, type(None)): return object_json elif isinstance(object_json, list): return [object_from_json(i) for i in object_json] elif isinstance(object_json, tuple): return tuple(object_from_json(i) for i in object_json) elif isinstance(object_json, dict): if "__type" not in object_json: # this is just a regular dictionary, e.g. the one in Parameter # containing parameterizations return {k: object_from_json(v) for k, v in object_json.items()} _type = object_json.pop("__type") if _type == "datetime": return datetime.datetime.strptime( object_json["value"], "%Y-%m-%d %H:%M:%S.%f" ) elif _type == "OrderedDict": return OrderedDict( [(k, object_from_json(v)) for k, v in object_json["value"]] ) elif _type == "DataFrame": return pd.read_json(object_json["value"]) elif _type not in DECODER_REGISTRY: err = ( f"The JSON dictionary passed to `object_from_json` has a type " f"{_type} that is not registered with a corresponding class in " f"DECODER_REGISTRY." ) raise JSONDecodeError(err) _class = DECODER_REGISTRY[_type] if issubclass(_class, enum.Enum): # to access enum members by name, use item access return _class[object_json["name"]] elif _class == GeneratorRun: return generator_run_from_json(object_json=object_json) elif _class == SimpleExperiment: return simple_experiment_from_json(object_json=object_json) elif _class == Experiment: return experiment_from_json(object_json=object_json) elif _class == SearchSpace: return search_space_from_json(search_space_json=object_json) return _class(**{k: object_from_json(v) for k, v in object_json.items()}) else: err = ( f"The object passed to `object_from_json` has an unsupported type: " f"{type(object_json)}." ) raise JSONDecodeError(err)
[docs]def generator_run_from_json(object_json: Dict[str, Any]) -> GeneratorRun: """Load Ax GeneratorRun from JSON.""" time_created_json = object_json.pop("time_created") type_json = object_json.pop("generator_run_type") index_json = object_json.pop("index") generator_run = GeneratorRun( **{k: object_from_json(v) for k, v in object_json.items()} ) generator_run._time_created = object_from_json(time_created_json) generator_run._generator_run_type = object_from_json(type_json) generator_run._index = object_from_json(index_json) return generator_run
[docs]def search_space_from_json(search_space_json: Dict[str, Any]) -> SearchSpace: """Load a SearchSpace from JSON. This function is necessary due to the coupled loading of SearchSpace and parameter constraints. """ parameters = object_from_json(search_space_json.pop("parameters")) json_param_constraints = search_space_json.pop("parameter_constraints") return SearchSpace( parameters=parameters, parameter_constraints=parameter_constraints_from_json( parameter_constraint_json=json_param_constraints, parameters=parameters ), )
[docs]def parameter_constraints_from_json( parameter_constraint_json: List[Dict[str, Any]], parameters: List[Parameter] ) -> List[ParameterConstraint]: """Load ParameterConstraints from JSON. Order and SumConstraint are tied to a search space, and require that SearchSpace's parameters to be passed in for decoding. Args: parameter_constraint_json: JSON representation of parameter constraints. parameters: Parameter definitions for decoding via parameter names. Returns: parameter_constraints: Python classes for parameter constraints. """ parameter_constraints = [] parameter_map = {p.name: p for p in parameters} for constraint in parameter_constraint_json: if constraint["__type"] == "OrderConstraint": lower_parameter = parameter_map[constraint["lower_name"]] upper_parameter = parameter_map[constraint["upper_name"]] parameter_constraints.append( OrderConstraint( lower_parameter=lower_parameter, upper_parameter=upper_parameter ) ) elif constraint["__type"] == "SumConstraint": parameters = [parameter_map[name] for name in constraint["parameter_names"]] parameter_constraints.append( SumConstraint( parameters=parameters, is_upper_bound=constraint["is_upper_bound"], bound=constraint["bound"], ) ) else: parameter_constraints.append(object_from_json(constraint)) return parameter_constraints
[docs]def trials_from_json( experiment: Experiment, trials_json: Dict[str, Any] ) -> Dict[int, BaseTrial]: """Load Ax Trials from JSON.""" loaded_trials = {} for index, batch_json in trials_json.items(): is_trial = batch_json["__type"] == "Trial" batch_json = { k: object_from_json(v) for k, v in batch_json.items() if k != "__type" } loaded_trials[int(index)] = ( trial_from_json(experiment=experiment, **batch_json) if is_trial else batch_trial_from_json(experiment=experiment, **batch_json) ) return loaded_trials
[docs]def data_from_json( data_by_trial_json: Dict[str, Any] ) -> Dict[int, "OrderedDict[int, Data]"]: """Load Ax Data from JSON.""" data_by_trial = object_from_json(data_by_trial_json) # hack necessary because Python's json module converts dictionary # keys to strings: https://stackoverflow.com/q/1450957 return { int(k): OrderedDict({int(k2): v2 for k2, v2 in v.items()}) for k, v in data_by_trial.items() }
[docs]def simple_experiment_from_json(object_json: Dict[str, Any]) -> SimpleExperiment: """Load AE SimpleExperiment from JSON.""" time_created_json = object_json.pop("time_created") trials_json = object_json.pop("trials") experiment_type_json = object_json.pop("experiment_type") data_by_trial_json = object_json.pop("data_by_trial") description_json = object_json.pop("description") is_test_json = object_json.pop("is_test") optimization_config = object_from_json(object_json.pop("optimization_config")) # not relevant to simple experiment del object_json["tracking_metrics"] del object_json["runner"] kwargs = {k: object_from_json(v) for k, v in object_json.items()} kwargs["evaluation_function"] = unimplemented_evaluation_function kwargs["objective_name"] = optimization_config.objective.metric.name kwargs["minimize"] = optimization_config.objective.minimize kwargs["outcome_constraints"] = optimization_config.outcome_constraints experiment = SimpleExperiment(**kwargs) experiment.description = object_from_json(description_json) experiment.is_test = object_from_json(is_test_json) experiment._time_created = object_from_json(time_created_json) experiment._trials = trials_from_json(experiment, trials_json) experiment._experiment_type = object_from_json(experiment_type_json) experiment._data_by_trial = data_from_json(data_by_trial_json) return experiment
[docs]def experiment_from_json(object_json: Dict[str, Any]) -> Experiment: """Load Ax Experiment from JSON.""" time_created_json = object_json.pop("time_created") trials_json = object_json.pop("trials") experiment_type_json = object_json.pop("experiment_type") data_by_trial_json = object_json.pop("data_by_trial") experiment = Experiment(**{k: object_from_json(v) for k, v in object_json.items()}) experiment._time_created = object_from_json(time_created_json) experiment._trials = trials_from_json(experiment, trials_json) for trial in experiment._trials.values(): for arm in trial.arms: experiment._arms_by_signature[arm.signature] = arm if experiment.status_quo is not None: sq_sig = experiment.status_quo.signature experiment._arms_by_signature[sq_sig] = experiment.status_quo experiment._experiment_type = object_from_json(experiment_type_json) experiment._data_by_trial = data_from_json(data_by_trial_json) return experiment