Example #1
0
    def new_trial(
        self,
        generator_run: Optional[GeneratorRun] = None,
        trial_type: Optional[str] = None,
        ttl_seconds: Optional[int] = None,
    ) -> Trial:
        """Create a new trial associated with this experiment.

        Args:
            generator_run: GeneratorRun, associated with this trial.
                Trial has only one generator run (and thus arm)
                attached to it. This can 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.
            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 Trial(
            experiment=self,
            trial_type=trial_type,
            generator_run=generator_run,
            ttl_seconds=ttl_seconds,
        )
Example #2
0
File: decoder.py Project: raff7/Ax
 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
Example #3
0
 def new_trial(
     self,
     generator_run: Optional[GeneratorRun] = None,
     trial_type: Optional[str] = None,
 ) -> Trial:
     """Create a new trial associated with this experiment."""
     return Trial(experiment=self,
                  trial_type=trial_type,
                  generator_run=generator_run)
Example #4
0
    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}",
            )
Example #5
0
    def new_trial(
        self,
        generator_run: Optional[GeneratorRun] = None,
        trial_type: Optional[str] = None,
    ) -> Trial:
        """Create a new trial associated with this experiment.

        Args:
            generator_run: GeneratorRun, associated with this trial.
                Trial has only one generator run (and thus arm)
                attached to it. This can 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.
        """
        return Trial(experiment=self,
                     trial_type=trial_type,
                     generator_run=generator_run)
Example #6
0
    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), 4)
        # Get them in the order we want for tests below
        if observations[0].features.parameters["x"] == 1:
            observations.reverse()

        obsd_truth = {
            "metric_names": [["a"], ["a"], ["b"], ["b"]],
            "means": [
                np.array([2.0]),
                np.array([3.0]),
                np.array([4.0]),
                np.array([5.0]),
            ],
            "covariance": [
                np.diag([4.0]),
                np.diag([9.0]),
                np.diag([16.0]),
                np.diag([25.0]),
            ],
        }
        cname_truth = ["0_0", "0_1", "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])
Example #7
0
    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"])
Example #8
0
    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
Example #9
0
def 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: GeneratorRun,
    runner: Optional[Runner],
    num_arms_created: int,
) -> Trial:
    """Load Ax trial from JSON.

    Other classes don't need explicit deserializers, because we can just use
    their constructors (see decoder.py). However, the constructor for Trial
    does not allow us to exactly recreate an existing object.
    """

    trial = Trial(experiment=experiment, generator_run=generator_run)
    trial._index = index
    trial._trial_type = trial_type
    trial._status = status
    trial._time_created = time_created
    trial._time_completed = time_completed
    trial._time_staged = time_staged
    trial._time_run_started = time_run_started
    trial._abandoned_reason = abandoned_reason
    trial._run_metadata = run_metadata or {}
    trial._runner = runner
    trial._num_arms_created = num_arms_created
    return trial
Example #10
0
def 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: GeneratorRun,
    runner: Optional[Runner],
    num_arms_created: int,
    # 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,
) -> Trial:
    """Load Ax trial from JSON.

    Other classes don't need explicit deserializers, because we can just use
    their constructors (see decoder.py). However, the constructor for Trial
    does not allow us to exactly recreate an existing object.
    """

    trial = Trial(
        experiment=experiment, generator_run=generator_run, ttl_seconds=ttl_seconds
    )
    trial._index = index
    trial._trial_type = trial_type
    # Swap `DISPATCHED` for `RUNNING`, since `DISPATCHED` is deprecated and nearly
    # equivalent to `RUNNING`.
    trial._status = status if status != TrialStatus.DISPATCHED else TrialStatus.RUNNING
    trial._time_created = time_created
    trial._time_completed = time_completed
    trial._time_staged = time_staged
    trial._time_run_started = time_run_started
    trial._abandoned_reason = abandoned_reason
    trial._run_metadata = run_metadata or {}
    trial._stop_metadata = stop_metadata or {}
    trial._runner = runner
    trial._num_arms_created = num_arms_created
    trial._generation_step_index = generation_step_index
    trial._properties = properties or {}
    warn_on_kwargs(callable_with_kwargs=Trial, **kwargs)
    return trial
Example #11
0
    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)
Example #12
0
    def testObservationsFromMapData(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",
                "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,
            },
            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",
                "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,
            },
            1: {
                "arm_name": "0_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,
            },
        }
        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", "z", "timestamp"]
        ]
        data = MapData(
            df=df,
            map_key_infos=[
                MapKeyInfo(key="z", default_value=0.0),
                MapKeyInfo(key="timestamp", default_value=0.0),
            ],
        )
        observations = observations_from_map_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"])
            self.assertEqual(obs.features.metadata, {"timestamp": t["timestamp"]})
Example #13
0
def 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: GeneratorRun,
    runner: Optional[Runner],
    num_arms_created: int,
    generation_step_index: Optional[int] = None,
) -> Trial:
    """Load Ax trial from JSON.

    Other classes don't need explicit deserializers, because we can just use
    their constructors (see decoder.py). However, the constructor for Trial
    does not allow us to exactly recreate an existing object.
    """

    trial = Trial(experiment=experiment, generator_run=generator_run)
    trial._index = index
    trial._trial_type = trial_type
    # Swap `DISPATCHED` for `RUNNING`, since `DISPATCHED` is deprecated and nearly
    # equivalent to `RUNNING`.
    trial._status = status if status != TrialStatus.DISPATCHED else TrialStatus.RUNNING
    trial._time_created = time_created
    trial._time_completed = time_completed
    trial._time_staged = time_staged
    trial._time_run_started = time_run_started
    trial._abandoned_reason = abandoned_reason
    trial._run_metadata = run_metadata or {}
    trial._runner = runner
    trial._num_arms_created = num_arms_created
    trial._generation_step_index = generation_step_index
    return trial