Fix Scheduler test_run_trials_and_yield_results_with_early_stopper th…

…at was timing out the CI (#2770)

Summary:
Pull Request resolved: #2770

The `EarlyStopsInsteadOfNormalCompletionScheduler` test class was overwriting `should_stop_trials_early`, which was being used to make the early stopping decisions prior to #2666. After that diff, the overwrite became a no-op, leading to this test hanging out and failing the CI.

Reviewed By: danielcohenlive, mpolson64

Differential Revision: D62784565

fbshipit-source-id: 11e3b6be0057e82afd82b706aeda54af686a7c8b

ax/service/tests/scheduler_test_utils.py

@@ -141,15 +141,6 @@ def report_results(self, force_refit: bool = False) -> dict[str, Any]:
         }
 
 
-class EarlyStopsInsteadOfNormalCompletionScheduler(TestScheduler):
-    """Test scheduler that marks all trials as ones that should be early-stopped."""
-
-    def should_stop_trials_early(
-        self, trial_indices: set[int]
-    ) -> dict[int, Optional[str]]:
-        return {i: None for i in trial_indices}
-
-
 # ---- Runners below simulate different usage and failure modes for scheduler ----
 
 
@@ -1310,7 +1301,7 @@ def _helper_for_run_trials_and_yield_results_with_early_stopper(
             experiment=self.branin_experiment,
             generation_strategy=self.two_sobol_steps_GS,
         )
-        scheduler = EarlyStopsInsteadOfNormalCompletionScheduler(
+        scheduler = TestScheduler(
             experiment=self.branin_experiment,  # Has runner and metrics.
             generation_strategy=gs,
             options=SchedulerOptions(
@@ -1319,10 +1310,9 @@ def _helper_for_run_trials_and_yield_results_with_early_stopper(
             db_settings=self.db_settings_if_always_needed,
         )
         # All trials should be marked complete after one run.
-        with patch.object(
-            scheduler,
-            "should_stop_trials_early",
-            wraps=scheduler.should_stop_trials_early,
+        with patch(
+            "ax.service.utils.early_stopping.should_stop_trials_early",
+            wraps=lambda trial_indices, **kwargs: {i: None for i in trial_indices},
         ) as mock_should_stop_trials_early, patch.object(
             InfinitePollRunner, "stop", return_value=None
         ) as mock_stop_trial_run: