def trial_from_sqa(self, trial_sqa: SQATrial, experiment: Experiment) -> BaseTrial: """Convert SQLAlchemy Trial to Ax Trial.""" if trial_sqa.is_batch: trial = BatchTrial(experiment=experiment, optimize_for_power=trial_sqa.optimize_for_power) generator_run_structs = [ GeneratorRunStruct( generator_run=self.generator_run_from_sqa( generator_run_sqa=generator_run_sqa), 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 else: new_generator_run_structs.append(struct) generator_run_structs = new_generator_run_structs trial._generator_run_structs = generator_run_structs 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 } else: trial = Trial(experiment=experiment) 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]) trial._index = trial_sqa.index 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]]`. trial._run_metadata = (dict(trial_sqa.run_metadata) if trial_sqa.run_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) return trial
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). """ 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, ), 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 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, ) 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 = ( # pyre-fixme[6]: Expected `Mapping[Variable[_KT], Variable[_VT]]` for # 1st param but got `Optional[Dict[str, typing.Any]]`. dict(trial_sqa.run_metadata) if trial_sqa.run_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 = trial_sqa.properties or {} trial.db_id = trial_sqa.id return trial