def new_batch_trial(
self,
generator_run: Optional[GeneratorRun] = None,
trial_type: Optional[str] = None,
optimize_for_power: Optional[bool] = False,
ttl_seconds: Optional[int] = None,
) -> BatchTrial:
"""Create a new batch trial associated with this experiment.
Args:
generator_run: GeneratorRun, associated with this trial. This can a
also be set later through `add_arm` or `add_generator_run`, but a
trial's associated generator run is immutable once set.
trial_type: Type of this trial, if used in MultiTypeExperiment.
optimize_for_power: Whether to optimize the weights of arms in this
trial such that the experiment's power to detect effects of
certain size is as high as possible. Refer to documentation of
`BatchTrial.set_status_quo_and_optimize_power` for more detail.
ttl_seconds: If specified, trials will be considered failed after
this many seconds since the time the trial was ran, unless the
trial is completed before then. Meant to be used to detect
'dead' trials, for which the evaluation process might have
crashed etc., and which should be considered failed after
their 'time to live' has passed.
"""
if ttl_seconds is not None:
self._trials_have_ttl = True
return BatchTrial(
experiment=self,
trial_type=trial_type,
generator_run=generator_run,
optimize_for_power=optimize_for_power,
ttl_seconds=ttl_seconds,
)
def new_batch_trial(
self,
generator_run: Optional[GeneratorRun] = None,
trial_type: Optional[str] = None,
) -> BatchTrial:
"""Create a new batch trial associated with this experiment."""
return BatchTrial(experiment=self,
trial_type=trial_type,
generator_run=generator_run)
def new_batch_trial(
self,
generator_run: Optional[GeneratorRun] = None,
trial_type: Optional[str] = None,
optimize_for_power: Optional[bool] = False,
) -> BatchTrial:
"""Create a new batch trial associated with this experiment."""
return BatchTrial(
experiment=self,
trial_type=trial_type,
generator_run=generator_run,
optimize_for_power=optimize_for_power,
)
def new_batch_trial(
self,
generator_run: Optional[GeneratorRun] = None,
trial_type: Optional[str] = None,
optimize_for_power: Optional[bool] = False,
) -> BatchTrial:
"""Create a new batch trial associated with this experiment.
Args:
generator_run: GeneratorRun, associated with this trial. This can a
also be set later through `add_arm` or `add_generator_run`, but a
trial's associated generator run is immutable once set.
trial_type: Type of this trial, if used in MultiTypeExperiment.
optimize_for_power: Whether to optimize the weights of arms in this
trial such that the experiment's power to detect effects of
certain size is as high as possible. Refer to documentation of
`BatchTrial.set_status_quo_and_optimize_power` for more detail.
"""
return BatchTrial(
experiment=self,
trial_type=trial_type,
generator_run=generator_run,
optimize_for_power=optimize_for_power,
)
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
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,
ttl_seconds=trial_sqa.ttl_seconds,
)
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, ttl_seconds=trial_sqa.ttl_seconds)
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 = (
# 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
return trial
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: float,
) -> 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)
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._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 = status_quo_weight
return batch
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
trial._status = trial_sqa.status
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 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 testObservationsFromDataAbandoned(self):
truth = {
0.5: {
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 2.0,
"sem": 2.0,
"trial_index": 0,
"metric_name": "a",
"updated_parameters": {"x": 0, "y": "a", "z": 0.5},
"mean_t": np.array([2.0]),
"covariance_t": np.array([[4.0]]),
"z": 0.5,
"timestamp": 50,
},
1: {
"arm_name": "1_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"updated_parameters": {"x": 0, "y": "a", "z": 1},
"mean_t": np.array([4.0]),
"covariance_t": np.array([[16.0]]),
"z": 1,
"timestamp": 100,
},
0.25: {
"arm_name": "2_0",
"parameters": {"x": 1, "y": "a", "z": 0.5},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"updated_parameters": {"x": 1, "y": "b", "z": 0.25},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
"z": 0.25,
"timestamp": 25,
},
0.75: {
"arm_name": "2_1",
"parameters": {"x": 1, "y": "b", "z": 0.75},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"updated_parameters": {"x": 1, "y": "b", "z": 0.75},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
"z": 0.75,
"timestamp": 25,
},
}
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"], parameters=obs["parameters"])
for _, obs in truth.items()
}
experiment = Mock()
experiment._trial_indices_by_status = {status: set() for status in TrialStatus}
trials = {
obs["trial_index"]: (
Trial(experiment, GeneratorRun(arms=[arms[obs["arm_name"]]]))
)
for _, obs in list(truth.items())[:-1]
if not obs["arm_name"].startswith("2")
}
batch = BatchTrial(experiment, GeneratorRun(arms=[arms["2_0"], arms["2_1"]]))
trials.update({2: batch})
trials.get(1).mark_abandoned()
trials.get(2).mark_arm_abandoned(arm_name="2_1")
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(list(truth.values()))[
["arm_name", "trial_index", "mean", "sem", "metric_name"]
]
data = Data(df=df)
# 1 arm is abandoned and 1 trial is abandoned, so only 2 observations should be
# included.
obs_no_abandoned = observations_from_data(experiment, data)
self.assertEqual(len(obs_no_abandoned), 2)
# 1 arm is abandoned and 1 trial is abandoned, so only 2 observations should be
# included.
obs_with_abandoned = observations_from_data(
experiment, data, include_abandoned=True
)
self.assertEqual(len(obs_with_abandoned), 4)