Refactor Ax service generator

* Simplify ingest function
* self._ax_client.attach_trial not called on ignored trial

Author: shuds13

optimas/generators/ax/service/base.py

@@ -23,7 +23,6 @@
     Trial,
     VaryingParameter,
     Parameter,
-    TrialStatus,
 )
 from optimas.generators.ax.base import AxGenerator
 from optimas.utils.ax import AxModelManager
@@ -157,6 +156,31 @@ def suggest(self, num_points: Optional[int]) -> List[dict]:
             points.append(point)
         return points
 
+    def ingest(self, results: List[dict]) -> None:
+        """Send the results of evaluations to the generator."""
+        for result in results:
+            trial = Trial.from_dict(
+                trial_dict=result,
+                varying_parameters=self._varying_parameters,
+                objectives=self._objectives,
+                analyzed_parameters=self._analyzed_parameters,
+                custom_parameters=self._custom_trial_parameters,
+            )
+            if trial.ignored:
+                continue
+            try:
+                ax_trial = self._ax_client.get_trial(trial.ax_trial_id)
+            except AttributeError:
+                ax_trial = self._insert_unknown_trial(trial)
+            finally:
+                if trial.completed:
+                    self._complete_trial(ax_trial.index, trial)
+                elif trial.failed:
+                    if self._abandon_failed_trials:
+                        ax_trial.mark_abandoned()
+                    else:
+                        ax_trial.mark_failed()
+
     def _ignore_out_of_bounds(self, trial: Trial) -> None:
         """Check if trial parameters are within their bounds."""
         for var, value in zip(trial.varying_parameters, trial.parameter_values):
@@ -177,83 +201,69 @@ def ignore_trials(self, trials: List[Trial]) -> None:
                 # Handle unknown trial
                 self._ignore_out_of_bounds(trial)
 
-    def ingest(self, results: List[dict]) -> None:
-        """Send the results of evaluations to the generator."""
-        for result in results:
-            trial = Trial.from_dict(
-                trial_dict=result,
-                varying_parameters=self._varying_parameters,
-                objectives=self._objectives,
-                analyzed_parameters=self._analyzed_parameters,
-                custom_parameters=self._custom_trial_parameters,
-            )
-            try:
-                trial_id = trial.ax_trial_id
-                ax_trial = self._ax_client.get_trial(trial_id)
-            except AttributeError:
-                params = {}
-                for var, value in zip(
-                    trial.varying_parameters, trial.parameter_values
-                ):
-                    params[var.name] = value
-                try:
-                    _, trial_id = self._ax_client.attach_trial(params)
-                except ValueError as error:
-                    # Bypass checks from AxClient and manually add a trial
-                    # outside of the search space.
-                    # https://github.com/facebook/Ax/issues/768#issuecomment-1036515242
-                    if "not a valid value" in str(error):
-                        if self._fit_out_of_design:
-                            ax_trial = self._ax_client.experiment.new_trial()
-                            ax_trial.add_arm(Arm(parameters=params))
-                            ax_trial.mark_running(no_runner_required=True)
-                            trial_id = ax_trial.index
-                    else:
-                        raise error
-                ax_trial = self._ax_client.get_trial(trial_id)
+    def _get_ingest_params(self, trial: Trial) -> Dict:
+        """Return a trials ingest parameters as a dictionary."""
+        params = {}
+        for var, value in zip(trial.varying_parameters, trial.parameter_values):
+            params[var.name] = value
+        return params
 
-                # Since data was given externally, reduce number of
-                # initialization trials, but only if they have not failed.
-                if trial.completed and not self._enforce_n_init:
-                    generation_strategy = self._ax_client.generation_strategy
-                    current_step = generation_strategy.current_step
-                    # Reduce only if there are still Sobol trials left.
-                    if current_step.model == Models.SOBOL:
-                        for tc in current_step.transition_criteria:
-                            # Looping over all criterial makes sure we reduce
-                            # the transition thresholds due to `_n_init`
-                            # (i.e., max trials) and `min_trials_observed=1` (
-                            # i.e., min trials).
-                            if isinstance(tc, (MinTrials, MaxTrials)):
-                                tc.threshold -= 1
-                        generation_strategy._maybe_transition_to_next_node()
-            finally:
-                if trial.ignored:
-                    continue
-                elif trial.completed:
-                    outcome_evals = {}
-                    # Add objective evaluations.
-                    for ev in trial.objective_evaluations:
-                        outcome_evals[ev.parameter.name] = (ev.value, ev.sem)
-                    # Add outcome constraints evaluations.
-                    ax_config = self._ax_client.experiment.optimization_config
-                    if ax_config.outcome_constraints:
-                        ocs = [
-                            oc.metric.name
-                            for oc in ax_config.outcome_constraints
-                        ]
-                        for ev in trial.parameter_evaluations:
-                            par_name = ev.parameter.name
-                            if par_name in ocs:
-                                outcome_evals[par_name] = (ev.value, ev.sem)
-                    self._ax_client.complete_trial(
-                        trial_index=trial_id, raw_data=outcome_evals
-                    )
-                elif trial.failed:
-                    if self._abandon_failed_trials:
-                        ax_trial.mark_abandoned()
-                    else:
-                        ax_trial.mark_failed()
+    def _insert_unknown_trial(self, trial: Trial) -> None:
+        """Insert an unknown trial into the Ax client."""
+        params = self._get_ingest_params(trial)
+        try:
+            _, trial_id = self._ax_client.attach_trial(params)
+        except ValueError as error:
+            # Bypass checks from AxClient and manually add a trial
+            # outside of the search space.
+            # https://github.com/facebook/Ax/issues/768#issuecomment-1036515242
+            if "not a valid value" in str(error):
+                if self._fit_out_of_design:
+                    ax_trial = self._ax_client.experiment.new_trial()
+                    ax_trial.add_arm(Arm(parameters=params))
+                    ax_trial.mark_running(no_runner_required=True)
+                    trial_id = ax_trial.index
+            else:
+                raise error
+        ax_trial = self._ax_client.get_trial(trial_id)
+
+        # Since data was given externally, reduce number of
+        # initialization trials, but only if they have not failed.
+        if trial.completed and not self._enforce_n_init:
+            generation_strategy = self._ax_client.generation_strategy
+            current_step = generation_strategy.current_step
+            # Reduce only if there are still Sobol trials left.
+            if current_step.model == Models.SOBOL:
+                for tc in current_step.transition_criteria:
+                    # Looping over all criterial makes sure we reduce
+                    # the transition thresholds due to `_n_init`
+                    # (i.e., max trials) and `min_trials_observed=1` (
+                    # i.e., min trials).
+                    if isinstance(tc, (MinTrials, MaxTrials)):
+                        tc.threshold -= 1
+                generation_strategy._maybe_transition_to_next_node()
+        return ax_trial
+
+    def _complete_trial(self, ax_trial_index: int, trial: Trial) -> None:
+        """Complete a trial."""
+        outcome_evals = {}
+        # Add objective evaluations.
+        for ev in trial.objective_evaluations:
+            outcome_evals[ev.parameter.name] = (ev.value, ev.sem)
+        # Add outcome constraints evaluations.
+        ax_config = self._ax_client.experiment.optimization_config
+        if ax_config.outcome_constraints:
+            ocs = [
+                oc.metric.name
+                for oc in ax_config.outcome_constraints
+            ]
+            for ev in trial.parameter_evaluations:
+                par_name = ev.parameter.name
+                if par_name in ocs:
+                    outcome_evals[par_name] = (ev.value, ev.sem)
+        self._ax_client.complete_trial(
+            trial_index=ax_trial_index, raw_data=outcome_evals
+        )
 
     def _create_ax_client(self) -> AxClient:
         """Create Ax client."""