def batch_trial_from_json( experiment: core.experiment.Experiment, index: int, trial_type: Optional[str], status: TrialStatus, time_created: datetime, time_completed: Optional[datetime], time_staged: Optional[datetime], time_run_started: Optional[datetime], abandoned_reason: Optional[str], run_metadata: Optional[Dict[str, Any]], generator_run_structs: List[GeneratorRunStruct], runner: Optional[Runner], abandoned_arms_metadata: Dict[str, AbandonedArm], num_arms_created: int, status_quo: Optional[Arm], status_quo_weight_override: float, optimize_for_power: Optional[bool], # Allowing default values for backwards compatibility with # objects stored before these fields were added. ttl_seconds: Optional[int] = None, generation_step_index: Optional[int] = None, properties: Optional[Dict[str, Any]] = None, stop_metadata: Optional[Dict[str, Any]] = None, **kwargs: Any, ) -> BatchTrial: """Load Ax BatchTrial from JSON. Other classes don't need explicit deserializers, because we can just use their constructors (see decoder.py). However, the constructor for Batch does not allow us to exactly recreate an existing object. """ batch = BatchTrial(experiment=experiment, ttl_seconds=ttl_seconds) batch._index = index batch._trial_type = trial_type batch._status = status batch._time_created = time_created batch._time_completed = time_completed batch._time_staged = time_staged batch._time_run_started = time_run_started batch._abandoned_reason = abandoned_reason batch._run_metadata = run_metadata or {} batch._stop_metadata = stop_metadata or {} batch._generator_run_structs = generator_run_structs batch._runner = runner batch._abandoned_arms_metadata = abandoned_arms_metadata batch._num_arms_created = num_arms_created batch._status_quo = status_quo batch._status_quo_weight_override = status_quo_weight_override batch.optimize_for_power = optimize_for_power batch._generation_step_index = generation_step_index batch._properties = properties batch._refresh_arms_by_name() # Trigger cache build warn_on_kwargs(callable_with_kwargs=BatchTrial, **kwargs) return batch
def trial_from_sqa(self, trial_sqa: SQATrial, experiment: Experiment, reduced_state: bool = False) -> BaseTrial: """Convert SQLAlchemy Trial to Ax Trial. Args: trial_sqa: `SQATrial` to decode. reduced_state: Whether to load trial's generator run(s) with a slightly reduced state (without model state, search space, and optimization config). """ immutable_ss_and_oc = experiment.immutable_search_space_and_opt_config if trial_sqa.is_batch: trial = BatchTrial( experiment=experiment, optimize_for_power=trial_sqa.optimize_for_power, ttl_seconds=trial_sqa.ttl_seconds, index=trial_sqa.index, ) generator_run_structs = [ GeneratorRunStruct( generator_run=self.generator_run_from_sqa( generator_run_sqa=generator_run_sqa, reduced_state=reduced_state, immutable_search_space_and_opt_config= immutable_ss_and_oc, ), weight=generator_run_sqa.weight or 1.0, ) for generator_run_sqa in trial_sqa.generator_runs ] if trial_sqa.status_quo_name is not None: new_generator_run_structs = [] for struct in generator_run_structs: if (struct.generator_run.generator_run_type == GeneratorRunType.STATUS_QUO.name): status_quo_weight = struct.generator_run.weights[0] trial._status_quo = struct.generator_run.arms[0] trial._status_quo_weight_override = status_quo_weight trial._status_quo_generator_run_db_id = ( struct.generator_run.db_id) trial._status_quo_arm_db_id = struct.generator_run.arms[ 0].db_id else: new_generator_run_structs.append(struct) generator_run_structs = new_generator_run_structs trial._generator_run_structs = generator_run_structs if not reduced_state: trial._abandoned_arms_metadata = { abandoned_arm_sqa.name: self.abandoned_arm_from_sqa( abandoned_arm_sqa=abandoned_arm_sqa) for abandoned_arm_sqa in trial_sqa.abandoned_arms } trial._refresh_arms_by_name() # Trigger cache build else: trial = Trial( experiment=experiment, ttl_seconds=trial_sqa.ttl_seconds, index=trial_sqa.index, ) if trial_sqa.generator_runs: if len(trial_sqa.generator_runs) != 1: raise SQADecodeError( # pragma: no cover "Cannot decode SQATrial to Trial because trial is not batched " "but has more than one generator run.") trial._generator_run = self.generator_run_from_sqa( generator_run_sqa=trial_sqa.generator_runs[0], reduced_state=reduced_state, immutable_search_space_and_opt_config=immutable_ss_and_oc, ) trial._trial_type = trial_sqa.trial_type # Swap `DISPATCHED` for `RUNNING`, since `DISPATCHED` is deprecated and nearly # equivalent to `RUNNING`. trial._status = (trial_sqa.status if trial_sqa.status != TrialStatus.DISPATCHED else TrialStatus.RUNNING) trial._time_created = trial_sqa.time_created trial._time_completed = trial_sqa.time_completed trial._time_staged = trial_sqa.time_staged trial._time_run_started = trial_sqa.time_run_started trial._abandoned_reason = trial_sqa.abandoned_reason # pyre-fixme[9]: _run_metadata has type `Dict[str, Any]`; used as # `Optional[Dict[str, Any]]`. # pyre-fixme[8]: Attribute has type `Dict[str, typing.Any]`; used as # `Optional[typing.Dict[Variable[_KT], Variable[_VT]]]`. trial._run_metadata = (dict(trial_sqa.run_metadata) if trial_sqa.run_metadata is not None else None) # pyre-fixme[9]: _run_metadata has type `Dict[str, Any]`; used as # `Optional[Dict[str, Any]]`. # pyre-fixme[8]: Attribute has type `Dict[str, typing.Any]`; used as # `Optional[typing.Dict[Variable[_KT], Variable[_VT]]]`. trial._stop_metadata = (dict(trial_sqa.stop_metadata) if trial_sqa.stop_metadata is not None else None) trial._num_arms_created = trial_sqa.num_arms_created trial._runner = (self.runner_from_sqa(trial_sqa.runner) if trial_sqa.runner else None) trial._generation_step_index = trial_sqa.generation_step_index trial._properties = dict(trial_sqa.properties or {}) trial.db_id = trial_sqa.id return trial