Extract early stopping replay utilities to OSS (#4744)

Summary:

Adds the `estimate_hypothetical_early_stopping_savings()` function to the OSS module. This function estimates potential compute savings by replaying an experiment with a default early stopping strategy.

Reviewed By: bernardbeckerman

Differential Revision: D90150341

Author: shrutipatel31

ax/early_stopping/experiment_replay.py

@@ -17,7 +17,10 @@
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.parameter import ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
+from ax.early_stopping.dispatch import get_default_ess_or_none
 from ax.early_stopping.strategies.base import BaseEarlyStoppingStrategy
+from ax.early_stopping.utils import estimate_early_stopping_savings
+from ax.exceptions.core import UnsupportedError
 from ax.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
@@ -29,6 +32,12 @@
 
 logger: Logger = get_logger(__name__)
 
+# Constants for experiment replay
+MAX_REPLAY_TRIALS: int = 50
+REPLAY_NUM_POINTS_PER_CURVE: int = 20
+MAX_PENDING_TRIALS: int = 5
+MIN_SAVINGS_THRESHOLD: float = 0.1  # 10% threshold
+
 
 def replay_experiment(
     historical_experiment: Experiment,
@@ -105,3 +114,55 @@ def replay_experiment(
     orchestrator.run_all_trials()
     logger.info(f"Replayed the experiment in {perf_counter() - start_time} seconds.")
     return experiment
+
+
+def estimate_hypothetical_early_stopping_savings(
+    experiment: Experiment,
+    metric: Metric,
+    max_pending_trials: int = MAX_PENDING_TRIALS,
+) -> float:
+    """Estimate hypothetical early stopping savings using experiment replay.
+
+    This function replays the experiment with a default early stopping strategy
+    to calculate what savings would have been achieved if early stopping were
+    enabled.
+
+    Args:
+        experiment: The experiment to analyze.
+        metric: The metric to use for early stopping replay.
+        max_pending_trials: Maximum number of pending trials for the replay
+            orchestrator. Defaults to 5.
+
+    Returns:
+        Estimated savings as a fraction (0.0 to 1.0).
+
+    Raises:
+        UnsupportedError: If early stopping savings cannot be estimated.
+            This can happen when:
+            - No default early stopping strategy is available for this experiment
+              (e.g., multi-objective, constrained, or non-MapMetric experiments)
+            - The experiment data does not have progression data for replay
+            - The experiment replay fails due to invalid experiment state
+    """
+    default_ess = get_default_ess_or_none(experiment=experiment)
+    if default_ess is None:
+        raise UnsupportedError(
+            "No default early stopping strategy available (multi-objective, "
+            "constrained, or non-MapMetric experiment)."
+        )
+
+    replayed_experiment = replay_experiment(
+        historical_experiment=experiment,
+        num_samples_per_curve=REPLAY_NUM_POINTS_PER_CURVE,
+        max_replay_trials=MAX_REPLAY_TRIALS,
+        metric=metric,
+        max_pending_trials=max_pending_trials,
+        early_stopping_strategy=default_ess,
+    )
+
+    if replayed_experiment is None:
+        raise UnsupportedError(
+            "Experiment data does not have progression data for replay."
+        )
+
+    return estimate_early_stopping_savings(experiment=replayed_experiment)

ax/early_stopping/tests/test_experiment_replay.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from unittest.mock import patch
+
+from ax.early_stopping.experiment_replay import (
+    estimate_hypothetical_early_stopping_savings,
+)
+from ax.exceptions.core import UnsupportedError
+from ax.utils.common.testutils import TestCase
+from ax.utils.testing.core_stubs import (
+    get_branin_experiment,
+    get_branin_experiment_with_timestamp_map_metric,
+)
+from pyre_extensions import none_throws
+
+
+class TestEstimateHypotheticalEss(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        # Experiment with MapMetric for tests that need a valid default ESS.
+        self.exp = get_branin_experiment_with_timestamp_map_metric()
+        self.metric = none_throws(self.exp.optimization_config).objective.metric
+
+    def test_estimate_hypothetical_ess_no_default_strategy(self) -> None:
+        """Test that UnsupportedError is raised when no default ESS is available."""
+        # Non-MapMetric experiment has no default ESS.
+        exp = get_branin_experiment(has_optimization_config=True)
+        metric = none_throws(exp.optimization_config).objective.metric
+
+        with self.assertRaises(UnsupportedError) as e:
+            estimate_hypothetical_early_stopping_savings(
+                experiment=exp,
+                metric=metric,
+            )
+
+        self.assertIn(
+            "No default early stopping strategy available",
+            str(e.exception),
+        )
+
+    def test_estimate_hypothetical_ess_no_progression_data(self) -> None:
+        """Test that UnsupportedError is raised when experiment has no progression
+        data."""
+        with patch(
+            "ax.early_stopping.experiment_replay.replay_experiment",
+            return_value=None,
+        ):
+            with self.assertRaises(UnsupportedError) as e:
+                estimate_hypothetical_early_stopping_savings(
+                    experiment=self.exp,
+                    metric=self.metric,
+                )
+
+            self.assertIn(
+                "Experiment data does not have progression data for replay",
+                str(e.exception),
+            )
+
+    def test_estimate_hypothetical_ess_success(self) -> None:
+        """Test that savings are returned when replay succeeds."""
+        with (
+            patch(
+                "ax.early_stopping.experiment_replay.replay_experiment",
+            ) as mock_replay,
+            patch(
+                "ax.early_stopping.experiment_replay.estimate_early_stopping_savings",
+                return_value=0.25,
+            ) as mock_estimate,
+        ):
+            result = estimate_hypothetical_early_stopping_savings(
+                experiment=self.exp,
+                metric=self.metric,
+            )
+
+            self.assertEqual(result, 0.25)
+            mock_replay.assert_called_once()
+            mock_estimate.assert_called_once()
+
+    def test_estimate_hypothetical_ess_exception(self) -> None:
+        """Test that exceptions from replay propagate to the caller."""
+        with patch(
+            "ax.early_stopping.experiment_replay.replay_experiment",
+            side_effect=ValueError("Experiment's name is None."),
+        ):
+            with self.assertRaises(ValueError) as e:
+                estimate_hypothetical_early_stopping_savings(
+                    experiment=self.exp,
+                    metric=self.metric,
+                )
+
+            self.assertIn("Experiment's name is None.", str(e.exception))