import pandas as pd
from ax.core.base_trial import BaseTrial, TrialStatus
from ax.core.data import Data
from ax.core.generator_run import GeneratorRun, GeneratorRunType
from ax.utils.common.testutils import TestCase
from ax.utils.testing.core_stubs import get_arms, get_experiment, get_objective
TEST_DATA = Data(
df=pd.DataFrame(
[
{
"arm_name": "0_0",
"metric_name": get_objective().metric.name,
"mean": 1.0,
"sem": 2.0,
"trial_index": 0,
}
]
)
)
class TrialTest(TestCase):
def setUp(self):
self.experiment = get_experiment()
self.trial = self.experiment.new_trial()
self.arm = get_arms()[0]
self.trial.add_arm(self.arm)
def setUp(self):
self.df = pd.DataFrame([
{
"arm_name": "0_0",
"mean": 2.0,
"sem": 0.2,
"trial_index": 1,
"metric_name": "a",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_0",
"mean": 1.8,
"sem": 0.3,
"trial_index": 1,
"metric_name": "b",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_1",
"mean": float("nan"),
"sem": float("nan"),
"trial_index": 1,
"metric_name": "a",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_1",
"mean": 3.7,
"sem": 0.5,
"trial_index": 1,
"metric_name": "b",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_2",
"mean": 0.5,
"sem": None,
"trial_index": 1,
"metric_name": "a",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_2",
"mean": float("nan"),
"sem": float("nan"),
"trial_index": 1,
"metric_name": "b",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
{
"arm_name": "0_2",
"mean": float("nan"),
"sem": float("nan"),
"trial_index": 1,
"metric_name": "c",
"start_time": "2018-01-01",
"end_time": "2018-01-02",
},
])
self.data = Data(df=self.df)
self.optimization_config = OptimizationConfig(
objective=Objective(metric=Metric(name="a")),
outcome_constraints=[
OutcomeConstraint(
metric=Metric(name="b"),
op=ComparisonOp.GEQ,
bound=0,
relative=False,
)
],
)
def complete_trial(
self,
trial_index: int,
# acceptable `raw_data` argument formats:
# 1) {metric_name -> (mean, standard error)}
# 2) (mean, standard error) and we assume metric name == objective name
# 3) only the mean, and we assume metric name == objective name and
# standard error == 0
raw_data: TEvaluationOutcome,
metadata: Optional[Dict[str, str]] = None,
sample_size: Optional[int] = None,
) -> None:
"""
Completes the trial with given metric values and adds optional metadata
to it.
Args:
trial_index: Index of trial within the experiment.
raw_data: Evaluation data for the trial. Can be a mapping from
metric name to a tuple of mean and SEM, just a tuple of mean and
SEM if only one metric in optimization, or just the mean if there
is no SEM.
metadata: Additional metadata to track about this run.
"""
assert isinstance(
trial_index, int
), f"Trial index must be an int, got: {trial_index}." # pragma: no cover
trial = self.experiment.trials[trial_index]
if not isinstance(trial, Trial):
raise NotImplementedError(
"Batch trial functionality is not yet available through Service API."
)
if metadata is not None:
trial._run_metadata = metadata
arm_name = not_none(trial.arm).name
if isinstance(raw_data, dict):
evaluations = {arm_name: raw_data}
elif isinstance(raw_data, tuple):
evaluations = {
arm_name: {
self.experiment.optimization_config.objective.metric.name:
raw_data
}
}
elif isinstance(raw_data, (float, int)):
evaluations = {
arm_name: {
self.experiment.optimization_config.objective.metric.name:
(
raw_data,
0.0,
)
}
}
elif isinstance(raw_data,
(np.float32, np.float64, np.int32, np.int64)):
evaluations = {
arm_name: {
self.experiment.optimization_config.objective.metric.name:
(
numpy_type_to_python_type(raw_data),
0.0,
)
}
}
else:
raise ValueError(
"Raw data has an invalid type. The data must either be in the form "
"of a dictionary of metric names to mean, sem tuples, "
"or a single mean, sem tuple, or a single mean.")
sample_sizes = {arm_name: sample_size} if sample_size else {}
data = Data.from_evaluations(evaluations, trial.index, sample_sizes)
trial.mark_completed()
self.experiment.attach_data(data)
self._updated_trials.append(trial_index)
self._save_experiment_if_possible()
def testObservationsWithCandidateMetadata(self):
SOME_METADATA_KEY = "metadatum"
truth = [
{
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a"
},
"mean": 2.0,
"sem": 2.0,
"trial_index": 0,
"metric_name": "a",
},
{
"arm_name": "1_0",
"parameters": {
"x": 1,
"y": "b"
},
"mean": 3.0,
"sem": 3.0,
"trial_index": 1,
"metric_name": "a",
},
]
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"],
parameters=obs["parameters"])
for obs in truth
}
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"]]],
candidate_metadata_by_arm_signature={
arms[obs["arm_name"]].signature: {
SOME_METADATA_KEY: f"value_{obs['trial_index']}"
}
},
),
)
for obs in truth
}
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(truth)[[
"arm_name", "trial_index", "mean", "sem", "metric_name"
]]
data = Data(df=df)
observations = observations_from_data(experiment, data)
for observation in observations:
self.assertEqual(
observation.features.metadata.get(SOME_METADATA_KEY),
f"value_{observation.features.trial_index}",
)
class TrialTest(TestCase):
def setUp(self):
self.experiment = get_experiment()
self.trial = self.experiment.new_trial()
self.arm = get_arms()[0]
self.trial.add_arm(self.arm)
def test_eq(self):
new_trial = self.experiment.new_trial()
self.assertNotEqual(self.trial, new_trial)
def test_basic_properties(self):
self.assertEqual(self.experiment, self.trial.experiment)
self.assertEqual(self.trial.index, 0)
self.assertEqual(self.trial.status, TrialStatus.CANDIDATE)
self.assertIsNotNone(self.trial.time_created)
self.assertEqual(self.trial.arms_by_name["0_0"], self.trial.arm)
self.assertEqual(self.trial.arms, [self.arm])
self.assertEqual(self.trial.abandoned_arms, [])
self.assertEqual(self.trial.generator_run.generator_run_type,
GeneratorRunType.MANUAL.name)
# Test empty arms
with self.assertRaises(AttributeError):
self.experiment.new_trial().arm_weights
def test_adding_new_trials(self):
new_arm = get_arms()[1]
new_trial = self.experiment.new_trial(generator_run=GeneratorRun(
arms=[new_arm]))
with self.assertRaises(ValueError):
self.experiment.new_trial(generator_run=GeneratorRun(
arms=get_arms()))
self.assertEqual(new_trial.arms_by_name["1_0"], new_arm)
with self.assertRaises(KeyError):
self.trial.arms_by_name["1_0"]
def test_add_trial_same_arm(self):
# Check that adding new arm w/out name works correctly.
new_trial1 = self.experiment.new_trial(generator_run=GeneratorRun(
arms=[self.arm.clone(clear_name=True)]))
self.assertEqual(new_trial1.arm.name, self.trial.arm.name)
self.assertFalse(new_trial1.arm is self.trial.arm)
# Check that adding new arm with name works correctly.
new_trial2 = self.experiment.new_trial(generator_run=GeneratorRun(
arms=[self.arm.clone()]))
self.assertEqual(new_trial2.arm.name, self.trial.arm.name)
self.assertFalse(new_trial2.arm is self.trial.arm)
arm_wrong_name = self.arm.clone(clear_name=True)
arm_wrong_name.name = "wrong_name"
with self.assertRaises(ValueError):
new_trial2 = self.experiment.new_trial(generator_run=GeneratorRun(
arms=[arm_wrong_name]))
def test_abandonment(self):
self.assertFalse(self.trial.is_abandoned)
self.trial.mark_abandoned(reason="testing")
self.assertTrue(self.trial.is_abandoned)
self.assertFalse(self.trial.status.is_failed)
@patch(
f"{BaseTrial.__module__}.{BaseTrial.__name__}.fetch_data",
return_value=TEST_DATA,
)
def test_objective_mean(self, _mock):
self.assertEqual(self.trial.objective_mean, 1.0)
@patch(f"{BaseTrial.__module__}.{BaseTrial.__name__}.fetch_data",
return_value=Data())
def test_objective_mean_empty_df(self, _mock):
self.assertIsNone(self.trial.objective_mean)
def testRepr(self):
repr_ = "Trial(experiment_name='test', index=0, status=TrialStatus.CANDIDATE)"
self.assertEqual(str(self.trial), repr_)
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)
def testObservationsFromDataWithSomeMissingTimes(self):
truth = [
{
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a"
},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
"start_time": 0,
},
{
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b"
},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"start_time": 0,
},
{
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a"
},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"start_time": None,
},
{
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b"
},
"mean": 5.0,
"sem": 5.0,
"trial_index": 2,
"metric_name": "b",
"start_time": None,
},
]
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"],
parameters=obs["parameters"])
for obs in truth
}
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 truth
}
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(truth)[[
"arm_name", "trial_index", "mean", "sem", "metric_name",
"start_time"
]]
data = Data(df=df)
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 2)
# Get them in the order we want for tests below
if observations[0].features.parameters["x"] == 1:
observations.reverse()
obsd_truth = {
"metric_names": [["a", "b"], ["a", "b"]],
"means": [np.array([2.0, 4.0]),
np.array([3.0, 5.0])],
"covariance": [np.diag([4.0, 16.0]),
np.diag([9.0, 25.0])],
}
cname_truth = ["0_0", "0_1"]
for i, obs in enumerate(observations):
self.assertEqual(obs.features.parameters, truth[i]["parameters"])
self.assertEqual(obs.features.trial_index, truth[i]["trial_index"])
self.assertEqual(obs.data.metric_names,
obsd_truth["metric_names"][i])
self.assertTrue(
np.array_equal(obs.data.means, obsd_truth["means"][i]))
self.assertTrue(
np.array_equal(obs.data.covariance,
obsd_truth["covariance"][i]))
self.assertEqual(obs.arm_name, cname_truth[i])
def testFetchAndStoreData(self):
n = 10
exp = self._setupBraninExperiment(n)
batch = exp.trials[0]
batch.mark_completed()
# Test fetch data
batch_data = batch.fetch_data()
self.assertEqual(len(batch_data.df), n)
exp_data = exp.fetch_data()
exp_data2 = exp.metrics["b"].fetch_experiment_data(exp)
self.assertEqual(len(exp_data2.df), 4 * n)
self.assertEqual(len(exp_data.df), 4 * n)
self.assertEqual(len(exp.arms_by_name), 4 * n)
# Verify that `metrics` kwarg to `experiment.fetch_data` is respected.
exp.add_tracking_metric(Metric(name="not_yet_on_experiment"))
exp.attach_data(
Data(df=pd.DataFrame.from_records([{
"arm_name": "0_0",
"metric_name": "not_yet_on_experiment",
"mean": 3,
"sem": 0,
"trial_index": 0,
}])))
self.assertEqual(
set(
exp.fetch_data(metrics=[Metric(
name="not_yet_on_experiment")]).df["metric_name"].values),
{"not_yet_on_experiment"},
)
# Verify data lookup is empty for trial that does not yet have data.
self.assertEqual(len(exp.lookup_data_for_trial(1)[0].df), 0)
# Test local storage
t1 = exp.attach_data(batch_data)
t2 = exp.attach_data(exp_data)
full_dict = exp.data_by_trial
self.assertEqual(len(full_dict), 2) # data for 2 trials
self.assertEqual(len(full_dict[0]), 3) # 3 data objs for batch 0
# Test retrieving original batch 0 data
self.assertEqual(len(exp.lookup_data_for_ts(t1).df), n)
self.assertEqual(len(exp.lookup_data_for_trial(0)[0].df), n)
# Test retrieving full exp data
self.assertEqual(len(exp.lookup_data_for_ts(t2).df), 4 * n)
with self.assertRaisesRegex(ValueError, ".* for metric"):
exp.attach_data(batch_data, combine_with_last_data=True)
new_data = Data(df=pd.DataFrame.from_records([{
"arm_name": "0_0",
"metric_name": "z",
"mean": 3,
"sem": 0,
"trial_index": 0,
}]))
t3 = exp.attach_data(new_data, combine_with_last_data=True)
self.assertEqual(len(full_dict[0]), 4) # 4 data objs for batch 0 now
self.assertIn("z",
exp.lookup_data_for_ts(t3).df["metric_name"].tolist())
# Verify we don't get the data if the trial is abandoned
batch._status = TrialStatus.ABANDONED
self.assertEqual(len(batch.fetch_data().df), 0)
self.assertEqual(len(exp.fetch_data().df), 3 * n)
# Verify we do get the stored data if there are an unimplemented metrics.
del exp._data_by_trial[0][
t3] # Remove attached data for nonexistent metric.
# Remove implemented metric that is `available_while_running`
# (and therefore not pulled from cache).
exp.remove_tracking_metric(metric_name="b")
exp.add_tracking_metric(Metric(name="b")) # Add unimplemented metric.
batch._status = TrialStatus.COMPLETED
# Data should be getting looked up now.
self.assertEqual(batch.fetch_data(), exp.lookup_data_for_ts(t1))
self.assertEqual(exp.fetch_data(), exp.lookup_data_for_ts(t1))
metrics_in_data = set(batch.fetch_data().df["metric_name"].values)
# Data for metric "z" should no longer be present since we removed it.
self.assertEqual(metrics_in_data, {"b"})
# Verify that `metrics` kwarg to `experiment.fetch_data` is respected
# when pulling looked-up data.
self.assertEqual(
exp.fetch_data(metrics=[Metric(name="not_on_experiment")]), Data())
def testObservationsFromDataWithFidelities(self):
truth = {
0.5: {
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
"fidelities": json.dumps({"z": 0.5}),
"updated_parameters": {
"x": 0,
"y": "a",
"z": 0.5
},
"mean_t": np.array([2.0]),
"covariance_t": np.array([[4.0]]),
},
0.25: {
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b",
"z": 0.5
},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"fidelities": json.dumps({"z": 0.25}),
"updated_parameters": {
"x": 1,
"y": "b",
"z": 0.25
},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
},
1: {
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"fidelities": json.dumps({"z": 1}),
"updated_parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean_t": np.array([4.0]),
"covariance_t": np.array([[16.0]]),
},
}
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 truth.items()
}
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",
"fidelities"
]]
data = Data(df=df)
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 3)
for obs in observations:
t = truth[obs.features.parameters["z"]]
self.assertEqual(obs.features.parameters, t["updated_parameters"])
self.assertEqual(obs.features.trial_index, t["trial_index"])
self.assertEqual(obs.data.metric_names, [t["metric_name"]])
self.assertTrue(np.array_equal(obs.data.means, t["mean_t"]))
self.assertTrue(
np.array_equal(obs.data.covariance, t["covariance_t"]))
self.assertEqual(obs.arm_name, t["arm_name"])
class TrialTest(TestCase):
def setUp(self):
self.experiment = get_experiment()
self.trial = self.experiment.new_trial()
self.arm = get_arms()[0]
self.trial.add_arm(self.arm)
def test_eq(self):
new_trial = self.experiment.new_trial()
self.assertNotEqual(self.trial, new_trial)
def test_basic_properties(self):
self.assertEqual(self.experiment, self.trial.experiment)
self.assertEqual(self.trial.index, 0)
self.assertEqual(self.trial.status, TrialStatus.CANDIDATE)
self.assertTrue(self.trial.status.is_candidate)
self.assertIsNotNone(self.trial.time_created)
self.assertEqual(self.trial.arms_by_name["0_0"], self.trial.arm)
self.assertEqual(self.trial.arms[0].signature, self.arm.signature)
self.assertEqual(self.trial.abandoned_arms, [])
self.assertEqual(
self.trial.generator_run.generator_run_type, GeneratorRunType.MANUAL.name
)
# Test empty arms
with self.assertRaises(AttributeError):
self.experiment.new_trial().arm_weights
self.trial._status = TrialStatus.RUNNING
self.assertTrue(self.trial.status.is_running)
self.trial._status = TrialStatus.COMPLETED
self.assertTrue(self.trial.status.is_completed)
self.assertTrue(self.trial.completed_successfully)
def test_adding_new_trials(self):
new_arm = get_arms()[1]
cand_metadata = {new_arm.signature: {"a": "b"}}
new_trial = self.experiment.new_trial(
generator_run=GeneratorRun(
arms=[new_arm], candidate_metadata_by_arm_signature=cand_metadata
)
)
with self.assertRaises(ValueError):
self.experiment.new_trial(generator_run=GeneratorRun(arms=get_arms()))
self.assertEqual(new_trial.arms_by_name["1_0"].signature, new_arm.signature)
with self.assertRaises(KeyError):
self.trial.arms_by_name["1_0"]
self.assertEqual(
new_trial._get_candidate_metadata_from_all_generator_runs(),
{"1_0": cand_metadata[new_arm.signature]},
)
self.assertEqual(
new_trial._get_candidate_metadata("1_0"), cand_metadata[new_arm.signature]
)
self.assertRaises(
ValueError, new_trial._get_candidate_metadata, "this_is_not_an_arm"
)
def test_add_trial_same_arm(self):
# Check that adding new arm w/out name works correctly.
new_trial1 = self.experiment.new_trial(
generator_run=GeneratorRun(arms=[self.arm.clone(clear_name=True)])
)
self.assertEqual(new_trial1.arm.name, self.trial.arm.name)
self.assertFalse(new_trial1.arm is self.trial.arm)
# Check that adding new arm with name works correctly.
new_trial2 = self.experiment.new_trial(
generator_run=GeneratorRun(arms=[self.arm.clone()])
)
self.assertEqual(new_trial2.arm.name, self.trial.arm.name)
self.assertFalse(new_trial2.arm is self.trial.arm)
arm_wrong_name = self.arm.clone(clear_name=True)
arm_wrong_name.name = "wrong_name"
with self.assertRaises(ValueError):
new_trial2 = self.experiment.new_trial(
generator_run=GeneratorRun(arms=[arm_wrong_name])
)
def test_abandonment(self):
self.assertFalse(self.trial.status.is_abandoned)
self.trial.mark_abandoned(reason="testing")
self.assertTrue(self.trial.status.is_abandoned)
self.assertFalse(self.trial.status.is_failed)
self.assertTrue(self.trial.did_not_complete)
def test_mark_as(self):
for terminal_status in (
TrialStatus.ABANDONED,
TrialStatus.FAILED,
TrialStatus.COMPLETED,
TrialStatus.EARLY_STOPPED,
):
self.setUp()
# Note: This only tests the no-runner case (and thus not staging)
for status in (TrialStatus.RUNNING, terminal_status):
kwargs = {}
if status == TrialStatus.RUNNING:
kwargs["no_runner_required"] = True
if status == TrialStatus.ABANDONED:
kwargs["reason"] = "test_reason"
self.trial.mark_as(status=status, **kwargs)
self.assertTrue(self.trial.status == status)
if status == TrialStatus.ABANDONED:
self.assertEqual(self.trial.abandoned_reason, "test_reason")
else:
self.assertIsNone(self.trial.abandoned_reason)
if status != TrialStatus.RUNNING:
self.assertTrue(self.trial.status.is_terminal)
if status in [
TrialStatus.RUNNING,
TrialStatus.EARLY_STOPPED,
TrialStatus.COMPLETED,
]:
self.assertTrue(self.trial.status.expecting_data)
else:
self.assertFalse(self.trial.status.expecting_data)
def test_stop(self):
# test bad old status
with self.assertRaisesRegex(ValueError, "Can only stop STAGED or RUNNING"):
self.trial.stop(new_status=TrialStatus.ABANDONED)
# test bad new status
self.trial.mark_running(no_runner_required=True)
with self.assertRaisesRegex(ValueError, "New status of a stopped trial must"):
self.trial.stop(new_status=TrialStatus.CANDIDATE)
# dummy runner for testing stopping functionality
class DummyStopRunner(Runner):
def run(self, trial):
pass
def stop(self, trial, reason):
return {"reason": reason} if reason else {}
# test valid stopping
for reason, new_status in itertools.product(
(None, "because"),
(TrialStatus.COMPLETED, TrialStatus.ABANDONED, TrialStatus.EARLY_STOPPED),
):
self.setUp()
self.trial._runner = DummyStopRunner()
self.trial.mark_running()
self.assertEqual(self.trial.status, TrialStatus.RUNNING)
self.trial.stop(new_status=new_status, reason=reason)
self.assertEqual(self.trial.status, new_status)
self.assertEqual(
self.trial.stop_metadata, {} if reason is None else {"reason": reason}
)
@patch(
f"{BaseTrial.__module__}.{BaseTrial.__name__}.lookup_data",
return_value=TEST_DATA,
)
def test_objective_mean(self, _mock):
self.assertEqual(self.trial.objective_mean, 1.0)
@patch(
f"{BaseTrial.__module__}.{BaseTrial.__name__}.fetch_data", return_value=Data()
)
def test_objective_mean_empty_df(self, _mock):
with self.assertRaisesRegex(ValueError, "No data was retrieved for trial"):
self.assertIsNone(self.trial.objective_mean)
def testRepr(self):
repr_ = (
"Trial(experiment_name='test', index=0, "
"status=TrialStatus.CANDIDATE, arm=Arm(name='0_0', "
"parameters={'w': 0.85, 'x': 1, 'y': 'baz', 'z': False}))"
)
self.assertEqual(str(self.trial), repr_)
def test_update_run_metadata(self):
self.assertEqual(len(self.trial.run_metadata), 0)
self.trial.update_run_metadata({"something": "new"})
self.assertEqual(self.trial.run_metadata, {"something": "new"})
def test_update_stop_metadata(self):
self.assertEqual(len(self.trial.stop_metadata), 0)
self.trial.update_stop_metadata({"something": "new"})
self.assertEqual(self.trial.stop_metadata, {"something": "new"})
def gen(
self,
experiment: Experiment,
new_data: Optional[Data] = None, # Take in just the new data.
n: int = 1,
**kwargs: Any,
) -> GeneratorRun:
"""Produce the next points in the experiment."""
self._set_experiment(experiment=experiment)
# Get arm signatures for each entry in new_data that is indeed new.
new_arms = self._get_new_arm_signatures(
experiment=experiment, new_data=new_data
)
enough_observed = (
len(self._observed) + len(new_arms)
) >= self._curr.min_arms_observed
unlimited_arms = self._curr.num_arms == -1
enough_generated = (
not unlimited_arms and len(self._generated) >= self._curr.num_arms
)
# Check that minimum observed_arms is satisfied if it's enforced.
if self._curr.enforce_num_arms and enough_generated and not enough_observed:
raise ValueError(
"All trials for current model have been generated, but not enough "
"data has been observed to fit next model. Try again when more data "
"are available."
)
# TODO[Lena, T44021164]: take into account failed trials. Potentially
# reduce `_generated` count when a trial mentioned in new data failed.
all_data = (
Data.from_multiple_data(data=[self._data, new_data])
if new_data
else self._data
)
if self._model is None:
# Instantiate the first model.
self._set_current_model(experiment=experiment, data=all_data)
elif enough_generated and enough_observed:
# Change to the next model.
self._change_model(experiment=experiment, data=all_data)
elif new_data is not None:
# We're sticking with the curr. model, but should update with new data.
# pyre-fixme[16]: `Optional` has no attribute `update`.
self._model.update(experiment=experiment, data=new_data)
kwargs = consolidate_kwargs(
kwargs_iterable=[self._curr.model_gen_kwargs, kwargs],
keywords=get_function_argument_names(not_none(self._model).gen),
)
gen_run = not_none(self._model).gen(n=n, **kwargs)
# If nothing failed, update known data, _generated, and _observed.
self._data = all_data
self._generated.extend([arm.signature for arm in gen_run.arms])
self._observed.extend(new_arms)
self._generator_runs.append(gen_run)
return gen_run
def get_observed_pareto_frontiers(
experiment: Experiment,
data: Optional[Data] = None,
rel: bool = True,
arm_names: Optional[List[str]] = None,
) -> List[ParetoFrontierResults]:
"""
Find all Pareto points from an experiment.
Uses only values as observed in the data; no modeling is involved. Makes no
assumption about the search space or types of parameters. If "data" is provided will
use that, otherwise will use all data attached to the experiment.
Uses all arms present in data; does not filter according to experiment
search space. If arm_names is specified, will filter to just those arm whose names
are given in the list.
Assumes experiment has a multiobjective optimization config from which the
objectives and outcome constraints will be extracted.
Will generate a ParetoFrontierResults for every pair of metrics in the experiment's
multiobjective optimization config.
Args:
experiment: The experiment.
data: Data to use for computing Pareto frontier. If not provided, will fetch
data from experiment.
rel: Relativize, if status quo on experiment.
arm_names: If provided, computes Pareto frontier only from among the provided
list of arm names.
Returns: ParetoFrontierResults that can be used with interact_pareto_frontier.
"""
if data is None:
data = experiment.fetch_data()
if experiment.optimization_config is None:
raise ValueError("Experiment must have an optimization config")
if arm_names is not None:
data = Data(data.df[data.df["arm_name"].isin(arm_names)])
mb = get_tensor_converter_model(experiment=experiment, data=data)
pareto_observations = observed_pareto_frontier(modelbridge=mb)
# Convert to ParetoFrontierResults
metric_names = [
metric.name for metric in
experiment.optimization_config.objective.metrics # pyre-ignore
]
pfr_means = {name: [] for name in metric_names}
pfr_sems = {name: [] for name in metric_names}
for obs in pareto_observations:
for i, name in enumerate(obs.data.metric_names):
pfr_means[name].append(obs.data.means[i])
pfr_sems[name].append(np.sqrt(obs.data.covariance[i, i]))
# Relativize as needed
if rel and experiment.status_quo is not None:
# Get status quo values
# pyre-fixme[16]: `Optional` has no attribute `name`.
sq_df = data.df[data.df["arm_name"] == experiment.status_quo.name]
sq_df = sq_df.to_dict(orient="list")
sq_means = {}
sq_sems = {}
# pyre-fixme[6]: Expected `_SupportsIndex` for 1st param but got `str`.
for i, metric in enumerate(sq_df["metric_name"]):
# pyre-fixme[6]: Expected `_SupportsIndex` for 1st param but got `str`.
sq_means[metric] = sq_df["mean"][i]
# pyre-fixme[6]: Expected `_SupportsIndex` for 1st param but got `str`.
sq_sems[metric] = sq_df["sem"][i]
# Relativize
for name in metric_names:
if np.isnan(sq_sems[name]) or np.isnan(pfr_sems[name]).any():
# Just relativize means
pfr_means[name] = [(mu / sq_means[name] - 1) * 100
for mu in pfr_means[name]]
else:
# Use delta method
pfr_means[name], pfr_sems[name] = relativize(
means_t=pfr_means[name],
sems_t=pfr_sems[name],
mean_c=sq_means[name],
sem_c=sq_sems[name],
as_percent=True,
)
absolute_metrics = []
else:
absolute_metrics = metric_names
objective_thresholds = {}
if experiment.optimization_config.objective_thresholds is not None: # pyre-ignore
for objth in experiment.optimization_config.objective_thresholds:
is_rel = objth.metric.name not in absolute_metrics
if objth.relative != is_rel:
raise ValueError(
f"Objective threshold for {objth.metric.name} has "
f"rel={objth.relative} but was specified here as rel={is_rel}"
)
objective_thresholds[objth.metric.name] = objth.bound
# Construct ParetoFrontResults for each pair
pfr_list = []
param_dicts = [obs.features.parameters for obs in pareto_observations]
pfr_arm_names = [obs.arm_name for obs in pareto_observations]
for metric_a, metric_b in combinations(metric_names, 2):
pfr_list.append(
ParetoFrontierResults(
param_dicts=param_dicts,
means=pfr_means,
sems=pfr_sems,
primary_metric=metric_a,
secondary_metric=metric_b,
absolute_metrics=absolute_metrics,
objective_thresholds=objective_thresholds,
arm_names=pfr_arm_names,
))
return pfr_list
def lookup_or_fetch_experiment_data_multi(
cls,
experiment: core.experiment.Experiment,
metrics: Iterable[Metric],
trials: Optional[Iterable[core.base_trial.BaseTrial]] = None,
**kwargs: Any,
) -> Data:
"""Fetch or lookup (with fallback to fetching) data for given metrics,
depending on whether they are available while running.
If metric is available while running, its data can change (and therefore
we should always re-fetch it). If metric is available only upon trial
completion, its data does not change, so we can look up that data on
the experiment and only fetch the data that is not already attached to
the experiment.
NOTE: If fetching data for a metrics class that is only available upon
trial completion, data fetched in this function (data that was not yet
available on experiment) will be attached to experiment.
"""
# If this metric is available while trial is running, just default to
# `fetch_experiment_data_multi`.
if cls.is_available_while_running():
return cls.fetch_experiment_data_multi(experiment=experiment,
metrics=metrics,
trials=trials,
**kwargs)
# If this metric is available only upon trial completion, look up data
# on experiment and only fetch data that is not already cached.
if trials is None:
completed_trials = experiment.trials_by_status[
core.base_trial.TrialStatus.COMPLETED]
else:
completed_trials = [t for t in trials if t.status.is_completed]
if not completed_trials:
return Data()
trials_data = []
for trial in completed_trials:
cached_trial_data = experiment.lookup_data_for_trial(
trial_index=trial.index)[0]
cached_metric_names = cached_trial_data.metric_names
metrics_to_fetch = [
m for m in metrics if m.name not in cached_metric_names
]
if not metrics_to_fetch:
# If all needed data fetched from cache, no need to fetch any other data
# for trial.
trials_data.append(cached_trial_data)
continue
try:
fetched_trial_data = cls.fetch_experiment_data_multi(
experiment=experiment,
metrics=metrics_to_fetch,
trials=[trial],
**kwargs,
)
except NotImplementedError:
# Metric does not implement fetching logic and only uses lookup.
fetched_trial_data = Data()
final_data = Data.from_multiple_data(
# pyre-fixme [6]: Incompatible paramtype: Expected `Data`
# but got `AbstractDataFrameData`.
[cached_trial_data, fetched_trial_data])
if not final_data.df.empty:
experiment.attach_data(final_data)
trials_data.append(final_data)
return Data.from_multiple_data(
trials_data, subset_metrics=[m.name for m in metrics])
def testEmptyData(self):
df = Data().df
self.assertTrue(df.empty)
self.assertTrue(set(df.columns == REQUIRED_COLUMNS))
self.assertTrue(Data.from_multiple_data([]).df.empty)
def testSetStatusQuo(self, mock_fit, mock_observations_from_data):
# NOTE: If empty data object is not passed, observations are not
# extracted, even with mock.
modelbridge = ModelBridge(
search_space=get_search_space_for_value(),
model=0,
experiment=get_experiment_for_value(),
data=Data(),
status_quo_name="1_1",
)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Alternatively, we can specify by features
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
get_experiment_for_value(),
0,
status_quo_features=get_observation1().features,
)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Alternatively, we can specify on experiment
# Put a dummy arm with SQ name 1_1 on the dummy experiment.
exp = get_experiment_for_value()
sq = Arm(name="1_1", parameters={"x": 3.0})
exp._status_quo = sq
# Check that we set SQ to arm 1_1
modelbridge = ModelBridge(get_search_space_for_value(), 0, [], exp, 0)
self.assertEqual(modelbridge.status_quo, get_observation1())
# Errors if features and name both specified
with self.assertRaises(ValueError):
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
exp,
0,
status_quo_features=get_observation1().features,
status_quo_name="1_1",
)
# Left as None if features or name don't exist
modelbridge = ModelBridge(get_search_space_for_value(),
0, [],
exp,
0,
status_quo_name="1_0")
self.assertIsNone(modelbridge.status_quo)
modelbridge = ModelBridge(
get_search_space_for_value(),
0,
[],
get_experiment_for_value(),
0,
status_quo_features=ObservationFeatures(parameters={
"x": 3.0,
"y": 10.0
}),
)
self.assertIsNone(modelbridge.status_quo)
def warm_start_from_old_experiment(
self,
old_experiment: Experiment,
copy_run_metadata: bool = False,
trial_statuses_to_copy: Optional[List[TrialStatus]] = None,
) -> List[Trial]:
"""Copy all completed trials with data from an old Ax expeirment to this one.
This function checks that the parameters of each trial are members of the
current experiment's search_space.
NOTE: Currently only handles experiments with 1-arm ``Trial``-s, not
``BatchTrial``-s as there has not yet been need for support of the latter.
Args:
old_experiment: The experiment from which to transfer trials and data
copy_run_metadata: whether to copy the run_metadata from the old experiment
trial_statuses_to_copy: All trials with a status in this list will be
copied. By default, copies all ``COMPLETED`` and ``ABANDONED`` trials.
Returns:
List of trials successfully copied from old_experiment to this one
"""
if len(self.trials) > 0:
raise ValueError( # pragma: no cover
f"Can only warm-start experiments that don't yet have trials. "
f"Experiment {self.name} has {len(self.trials)} trials.")
old_parameter_names = set(
old_experiment.search_space.parameters.keys())
parameter_names = set(self.search_space.parameters.keys())
if old_parameter_names.symmetric_difference(parameter_names):
raise ValueError( # pragma: no cover
f"Cannot warm-start experiment '{self.name}' from experiment "
f"'{old_experiment.name}' due to mismatch in search space parameters."
f"Parameters in '{self.name}' but not in '{old_experiment.name}': "
f"{old_parameter_names - parameter_names}. Vice-versa: "
f"{parameter_names - old_parameter_names}.")
trial_statuses_to_copy = (
trial_statuses_to_copy if trial_statuses_to_copy is not None else
[TrialStatus.COMPLETED, TrialStatus.ABANDONED])
warm_start_trials = [
trial for trial in old_experiment.trials.values()
if trial.status in trial_statuses_to_copy
]
copied_trials = []
for trial in warm_start_trials:
if not isinstance(trial, Trial):
raise NotImplementedError( # pragma: no cover
"Only experiments with 1-arm trials currently supported.")
self.search_space.check_membership(not_none(trial.arm).parameters,
raise_error=True)
dat, ts = old_experiment.lookup_data_for_trial(
trial_index=trial.index)
is_completed_with_data = (trial.status == TrialStatus.COMPLETED
and ts != -1 and not dat.df.empty)
if is_completed_with_data or trial.status == TrialStatus.ABANDONED:
# Set trial index and arm name to their values in new trial.
new_trial = self.new_trial()
new_trial.add_arm(not_none(trial.arm).clone(clear_name=True))
new_trial.mark_running(no_runner_required=True)
new_trial.update_run_metadata(
{"run_id": trial.run_metadata.get("run_id")})
new_trial._properties["source"] = (
f"Warm start from Experiment: `{old_experiment.name}`, "
f"trial: `{trial.index}`")
if copy_run_metadata:
new_trial._run_metadata = trial.run_metadata
# Trial has data, so we replicate it on the new experiment.
if is_completed_with_data:
new_df = dat.df.copy()
new_df["trial_index"].replace(
{trial.index: new_trial.index}, inplace=True)
new_df["arm_name"].replace(
{
not_none(trial.arm).name: not_none(
new_trial.arm).name
},
inplace=True,
)
# Attach updated data to new trial on experiment and mark trial
# as completed.
self.attach_data(data=Data(df=new_df))
new_trial.mark_completed()
else:
new_trial.mark_abandoned(reason=trial.abandoned_reason)
copied_trials.append(new_trial)
if self._name is not None:
logger.info(
f"Copied {len(copied_trials)} completed trials and their data "
f"from {old_experiment.name} to {self.name}.")
else:
logger.info(
f"Copied {len(copied_trials)} completed trials and their data "
f"from {old_experiment.name}.")
return copied_trials
