Move Ax service gen to use _id

Author: shuds13

optimas/core/trial.py

@@ -63,6 +63,7 @@ def __init__(
         )
         evaluations = [] if evaluations is None else evaluations
         self._index = index
+        self._gen_id = None
         self._custom_parameters = (
             [] if custom_parameters is None else custom_parameters
         )
@@ -87,14 +88,14 @@ def to_dict(self) -> Dict:
             **self.objectives_as_dict(),
             **self.analyzed_parameters_as_dict(),
             **self.custom_parameters_as_dict(),
-            "_id": self._index,
+            "_id": self._gen_id,
             "_ignored": self._ignored,
             "_ignored_reason": self._ignored_reason,
             "_status": self._status,
         }
 
-        if hasattr(self, "_ax_trial_id"):
-            trial_dict["ax_trial_id"] = self.ax_trial_id
+        if hasattr(self, "_gen_id"):
+            trial_dict["gen_id"] = self.gen_id
 
         return trial_dict
 
@@ -141,9 +142,7 @@ def from_dict(
             custom_parameters=custom_parameters,
         )
         if "_id" in trial_dict:
-            trial._index = trial_dict["_id"]
-        if "ax_trial_id" in trial_dict:
-            trial._ax_trial_id = trial_dict["ax_trial_id"]
+            trial._gen_id = trial_dict["_id"]
         if "_ignored" in trial_dict:
             trial._ignored = trial_dict["_ignored"]
         if "_ignored_reason" in trial_dict:
@@ -204,13 +203,13 @@ def index(self, value):
         self._index = value
 
     @property
-    def ax_trial_id(self) -> int:
+    def gen_id(self) -> int:
         """Get the index of the trial."""
-        return self._ax_trial_id
+        return self._gen_id
 
-    @ax_trial_id.setter
-    def ax_trial_id(self, value):
-        self._ax_trial_id = value
+    @gen_id.setter
+    def gen_id(self, value):
+        self._gen_id = value
 
     @property
     def ignored(self) -> bool:

optimas/generators/ax/service/base.py

@@ -107,7 +107,7 @@ def __init__(
         save_model: Optional[bool] = True,
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
-    ) -> None:
+    ) -> None:       
         super().__init__(
             varying_parameters=varying_parameters,
             objectives=objectives,
@@ -152,7 +152,7 @@ def suggest(self, num_points: Optional[int]) -> List[dict]:
                 var.name: parameters.get(var.name)
                 for var in self._varying_parameters
             }
-            point["ax_trial_id"] = trial_id
+            point["_id"] = trial_id
             points.append(point)
         return points
 
@@ -169,8 +169,10 @@ def ingest(self, results: List[dict]) -> None:
             if trial.ignored:
                 continue
             try:
-                ax_trial = self._ax_client.get_trial(trial.ax_trial_id)
-            except AttributeError:
+                ax_trial = self._ax_client.get_trial(trial.gen_id)
+            except KeyError:
+                # This could indicate gen_id is not set or it is unrecognized
+                # Either way, Ax should generate a new trial / internal id.
                 ax_trial = self._insert_unknown_trial(trial)
             finally:
                 if trial.completed:
@@ -197,7 +199,7 @@ def _ignore_out_of_bounds(self, trial: Trial) -> None:
     def ignore_trials(self, trials: List[Trial]) -> None:
         """Ignore trials as determined by the generator."""
         for trial in trials:
-            if not hasattr(trial, "ax_trial_id"):
+            if trial.gen_id is None:
                 # Handle unknown trial
                 self._ignore_out_of_bounds(trial)
 
@@ -345,7 +347,7 @@ def _update_parameter(self, parameter):
 
     def _mark_trial_as_failed(self, trial: Trial):
         """Mark a trial as failed so that is not used for fitting the model."""
-        ax_trial = self._ax_client.get_trial(trial.ax_trial_id)
+        ax_trial = self._ax_client.get_trial(trial.gen_id)
         if self._abandon_failed_trials:
             ax_trial.mark_abandoned(unsafe=True)
         else:

optimas/generators/base.py

@@ -216,8 +216,6 @@ def ask_trials(self, n_trials: int) -> List[Trial]:
                     analyzed_parameters=self._analyzed_parameters,
                     custom_parameters=self._custom_trial_parameters,
                 )
-                if "ax_trial_id" in point:
-                    trial.ax_trial_id = point["ax_trial_id"]
                 gen_trials.append(trial)
             # Keep only trials that have been given data.
             for trial in gen_trials: