From 7c6e286a40a2c03f80ccc1c81da70eee6bbd4796 Mon Sep 17 00:00:00 2001
From: Lena Kashtelyan <drfreund@meta.com>
Date: Fri, 16 Aug 2024 15:06:31 -0700
Subject: [PATCH] Remove unnecessary wrapper-helper around early-stopping call
 (#2666)

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

As titled

Differential Revision: D55710758
---
 ax/service/scheduler.py | 22 +++-------------------
 1 file changed, 3 insertions(+), 19 deletions(-)

diff --git a/ax/service/scheduler.py b/ax/service/scheduler.py
index 38473787924..54a8a4623bb 100644
--- a/ax/service/scheduler.py
+++ b/ax/service/scheduler.py
@@ -1229,8 +1229,10 @@ def poll_and_process_results(self, poll_all_trial_statuses: bool = False) -> boo
         trial_indices_with_updated_data_or_status.update(trial_indices_with_new_data)
 
         # EARLY STOP TRIALS
-        stop_trial_info = self.should_stop_trials_early(
+        stop_trial_info = early_stopping_utils.should_stop_trials_early(
+            early_stopping_strategy=self.options.early_stopping_strategy,
             trial_indices=self.experiment.running_trial_indices,
+            experiment=self.experiment,
         )
         self.stop_trial_runs(
             trials=[self.experiment.trials[trial_idx] for trial_idx in stop_trial_info],
@@ -1424,24 +1426,6 @@ def _process_completed_trials(self, newly_completed: set[int]) -> None:
             trial_indices=newly_completed,
         )
 
-    def should_stop_trials_early(
-        self, trial_indices: set[int]
-    ) -> dict[int, Optional[str]]:
-        """Evaluate whether to early-stop running trials.
-
-        Args:
-            trial_indices: Indices of trials to consider for early stopping.
-
-        Returns:
-            A dictionary mapping trial indices that should be early stopped to
-            (optional) messages with the associated reason.
-        """
-        return early_stopping_utils.should_stop_trials_early(
-            early_stopping_strategy=self.options.early_stopping_strategy,
-            trial_indices=trial_indices,
-            experiment=self.experiment,
-        )
-
     def estimate_early_stopping_savings(self, map_key: Optional[str] = None) -> float:
         """Estimate early stopping savings using progressions of the MapMetric present
         on the EarlyStoppingConfig as a proxy for resource usage.