From 90a28def92d0beff78b5027c9a806885ddf190a4 Mon Sep 17 00:00:00 2001
From: Daniel Cohen <danielcohennyc@meta.com>
Date: Thu, 3 Oct 2024 12:53:11 -0700
Subject: [PATCH] Only generate PTS candidates if running trials have enough
 data (#2820)

Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/2820

Differential Revision: D63788500
---
 ax/service/scheduler.py                  |  1 -
 ax/service/tests/scheduler_test_utils.py | 56 +-----------------------
 ax/service/tests/test_scheduler.py       |  4 --
 3 files changed, 2 insertions(+), 59 deletions(-)

diff --git a/ax/service/scheduler.py b/ax/service/scheduler.py
index 46360925975..b0d9f4be588 100644
--- a/ax/service/scheduler.py
+++ b/ax/service/scheduler.py
@@ -1828,7 +1828,6 @@ def generate_candidates(
         Returns:
             List of trials, empty if generation is not possible.
         """
-        self.poll_and_process_results()
         new_trials = self._get_next_trials(
             num_trials=num_trials,
             n=self.options.batch_size or 1,
diff --git a/ax/service/tests/scheduler_test_utils.py b/ax/service/tests/scheduler_test_utils.py
index 77af969c44f..471b2d2ac63 100644
--- a/ax/service/tests/scheduler_test_utils.py
+++ b/ax/service/tests/scheduler_test_utils.py
@@ -80,7 +80,6 @@
     get_branin_search_space,
     get_generator_run,
     get_online_sobol_gpei_generation_strategy,
-    get_outcome_constraint,
     get_sobol,
     SpecialGenerationStrategy,
 )
@@ -2504,11 +2503,9 @@ def test_generate_candidates_works_for_iteration(self) -> None:
             ),
             db_settings=self.db_settings_if_always_needed,
         )
-        # AND GIVEN a sobol trial is running
+        # AND GIVEN a sobol trial is running with data
         scheduler.run(max_new_trials=1)
-        # if there is already data, the test doesn't prove that
-        # `generate_candidates()` fetches
-        self.assertTrue(scheduler.experiment.lookup_data().df.empty)
+        scheduler.poll_and_process_results()
 
         # WHEN generating candidates
         scheduler.generate_candidates(num_trials=1)
@@ -2609,55 +2606,6 @@ def test_generate_candidates_does_not_generate_if_missing_opt_config(self) -> No
         # THEN the experiment should have not generated candidates
         self.assertEqual(len(scheduler.experiment.trials), 1)
 
-    def test_generate_candidates_does_not_generate_if_overconstrained(self) -> None:
-        # GIVEN a scheduler using a GS with GPEI
-        gs = self._get_generation_strategy_strategy_for_test(
-            experiment=self.branin_experiment,
-            generation_strategy=get_online_sobol_gpei_generation_strategy(),
-        )
-        # this is a HITL experiment, so we don't want trials completing on their own.
-        self.branin_experiment.runner = InfinitePollRunner()
-        scheduler = Scheduler(
-            experiment=self.branin_experiment,
-            generation_strategy=gs,
-            options=SchedulerOptions(
-                init_seconds_between_polls=0,  # No wait bw polls so test is fast.
-                batch_size=10,
-                trial_type=TrialType.BATCH_TRIAL,
-                **self.scheduler_options_kwargs,
-            ),
-            db_settings=self.db_settings_if_always_needed,
-        )
-        # AND GIVEN a sobol trial is running
-        scheduler.run(max_new_trials=1)
-        # assert `run()` worked
-        self.assertEqual(len(scheduler.experiment.running_trial_indices), 1)
-        # AND GIVEN the optimization config is overconstrained
-        self.branin_experiment.optimization_config.outcome_constraints = [
-            get_outcome_constraint(
-                metric=get_branin_metric(name="branin_constraint"),
-                bound=20.0,
-                relative=True,
-            )
-        ]
-        if self.ALWAYS_USE_DB:
-            save_experiment(self.branin_experiment, config=self.db_config)
-        self.assertTrue(scheduler.experiment.lookup_data().df.empty)
-
-        # WHEN generating candidates
-        scheduler.generate_candidates(num_trials=1)
-
-        # THEN the experiment should have a GPEI generated trial
-        self.assertFalse(scheduler.experiment.lookup_data().df.empty)
-        self.assertEqual(
-            len(scheduler.experiment.trials), 1, str(scheduler.experiment.trials)
-        )
-        self.assertEqual(
-            len(scheduler.experiment.running_trial_indices),
-            1,
-            str(scheduler.experiment.trials),
-        )
-
     def test_compute_analyses(self) -> None:
         scheduler = Scheduler(
             experiment=self.branin_experiment,
diff --git a/ax/service/tests/test_scheduler.py b/ax/service/tests/test_scheduler.py
index 33c5808eb73..eeaee3d10af 100644
--- a/ax/service/tests/test_scheduler.py
+++ b/ax/service/tests/test_scheduler.py
@@ -145,10 +145,6 @@ def test_failure_rate_all_failed(self) -> None:
         self.branin_experiment.update_runner("type1", RunnerWithAllFailedTrials())
         super().test_failure_rate_all_failed()
 
-    def test_generate_candidates_does_not_generate_if_overconstrained(self) -> None:
-        self.branin_experiment.update_runner("type1", InfinitePollRunner())
-        super().test_generate_candidates_does_not_generate_if_overconstrained()
-
     def test_run_trials_and_yield_results_with_early_stopper(self) -> None:
         self.branin_experiment.update_runner("type1", InfinitePollRunner())
         super().test_run_trials_and_yield_results_with_early_stopper()