Merge pull request #207 from optimas-org/bug/fix_n_init

Fix crash when `n_init < sim_workers`
optimas-org · May 23, 2024 · da8b1f8 · da8b1f8
2 parents bb4ab72 + 193b2a3
commit da8b1f8
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 10 deletions.
diff --git a/optimas/generators/ax/service/base.py b/optimas/generators/ax/service/base.py
@@ -15,6 +15,7 @@
     GenerationStep,
     GenerationStrategy,
 )
+from ax.exceptions.core import DataRequiredError
 
 from optimas.core import (
     Objective,
@@ -238,6 +239,22 @@ def _create_ax_objectives(self) -> Dict[str, ObjectiveProperties]:
             objectives[obj.name] = ObjectiveProperties(minimize=obj.minimize)
         return objectives
 
+    def _create_sobol_step(self) -> GenerationStep:
+        """Create a Sobol generation step with `n_init` trials."""
+        # Ensure that at least 1 trial is completed before moving onto the BO
+        # step, and keep generating Sobol trials until that happens, even if
+        # the number of Sobol trials exceeds `n_init`.
+        # Otherwise, if we move to the BO step before any trial is completed,
+        # the next `ask` would fail with a `DataRequiredError`.
+        # This also allows the generator to work well when
+        # `sim_workers` > `n_init`.
+        return GenerationStep(
+            model=Models.SOBOL,
+            num_trials=self._n_init,
+            min_trials_observed=1,
+            enforce_num_trials=False,
+        )
+
     def _create_generation_steps(
         self, bo_model_kwargs: Dict
     ) -> List[GenerationStep]:

diff --git a/optimas/generators/ax/service/multi_fidelity.py b/optimas/generators/ax/service/multi_fidelity.py
@@ -117,9 +117,7 @@ def _create_generation_steps(
         steps = []
 
         # Add Sobol initialization with `n_init` random trials.
-        steps.append(
-            GenerationStep(model=Models.SOBOL, num_trials=self._n_init)
-        )
+        steps.append(self._create_sobol_step())
 
         # Continue indefinitely with GPKG.
         steps.append(

diff --git a/optimas/generators/ax/service/single_fidelity.py b/optimas/generators/ax/service/single_fidelity.py
@@ -156,9 +156,7 @@ def _create_generation_steps(
         steps = []
 
         # Add Sobol initialization with `n_init` random trials.
-        steps.append(
-            GenerationStep(model=Models.SOBOL, num_trials=self._n_init)
-        )
+        steps.append(self._create_sobol_step())
 
         # Continue indefinitely with BO.
         steps.append(

diff --git a/tests/test_ax_generators.py b/tests/test_ax_generators.py
@@ -692,8 +692,8 @@ def test_ax_service_init():
     var2 = VaryingParameter("x1", -5.0, 15.0)
     obj = Objective("f", minimize=False)
 
-    n_init = 4
-    n_external = 6
+    n_init = 2
+    n_external = 4
 
     for i in range(n_external):
         gen = AxSingleFidelityGenerator(
@@ -704,7 +704,7 @@ def test_ax_service_init():
             generator=gen,
             evaluator=ev,
             max_evals=6,
-            sim_workers=2,
+            sim_workers=3,  # Test with sim_workers > n_init.
             exploration_dir_path=f"./tests_output/test_ax_service_init_{i}",
         )
 
@@ -753,7 +753,7 @@ def test_ax_service_init():
         generator=gen,
         evaluator=ev,
         max_evals=15,
-        sim_workers=2,
+        sim_workers=3,  # Test with sim_workers > n_init.
         exploration_dir_path="./tests_output/test_ax_service_init_enforce",
     )