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.

Differential Revision: D90150341

Author: shrutipatel31

ax/early_stopping/experiment_replay.py

@@ -17,7 +17,9 @@
 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.generation_strategy.generation_strategy import (
     GenerationStep,
     GenerationStrategy,
@@ -29,6 +31,11 @@
 
 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
+
 
 def replay_experiment(
     historical_experiment: Experiment,
@@ -105,3 +112,56 @@ 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 | None:
+    """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), or None if:
+        - No default early stopping strategy is available for this experiment
+          (e.g., multi-objective, constrained, or non-MapMetric experiments)
+        - The experiment replay failed
+    """
+    try:
+        default_ess = get_default_ess_or_none(experiment=experiment)
+        if default_ess is None:
+            logger.info(
+                "No default early stopping strategy available (multi-objective, "
+                "constrained, or non-MapMetric experiment)."
+            )
+            return None
+
+        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:
+            logger.info("Experiment data does not have progression data for replay.")
+            return None
+
+        return estimate_early_stopping_savings(experiment=replayed_experiment)
+    except Exception as e:
+        # Replay can fail due to invalid experiment state (e.g., missing name,
+        # incompatible data format) or internal errors during orchestration.
+        logger.info(f"Experiment replay failed with exception: {e}")
+        return None

ax/early_stopping/tests/test_experiment_replay.py

@@ -0,0 +1,105 @@
+#!/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,
+    logger,
+)
+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 None is returned 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 patch.object(logger, "info") as mock_info:
+            result = estimate_hypothetical_early_stopping_savings(
+                experiment=exp,
+                metric=metric,
+            )
+
+            self.assertIsNone(result)
+            mock_info.assert_called_once()
+            self.assertIn(
+                "No default early stopping strategy available",
+                mock_info.call_args[0][0],
+            )
+
+    def test_estimate_hypothetical_ess_no_progression_data(self) -> None:
+        """Test that None is returned when experiment has no progression data."""
+        with (
+            patch(
+                "ax.early_stopping.experiment_replay.replay_experiment",
+                return_value=None,
+            ),
+            patch.object(logger, "info") as mock_info,
+        ):
+            result = estimate_hypothetical_early_stopping_savings(
+                experiment=self.exp,
+                metric=self.metric,
+            )
+
+            self.assertIsNone(result)
+            mock_info.assert_called_once_with(
+                "Experiment data does not have progression data for replay."
+            )
+
+    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 None is returned when replay raises an exception."""
+        with (
+            patch(
+                "ax.early_stopping.experiment_replay.replay_experiment",
+                side_effect=ValueError("Experiment's name is None."),
+            ),
+            patch.object(logger, "info") as mock_info,
+        ):
+            result = estimate_hypothetical_early_stopping_savings(
+                experiment=self.exp,
+                metric=self.metric,
+            )
+
+            self.assertIsNone(result)
+            mock_info.assert_called_once_with(
+                "Experiment replay failed with exception: Experiment's name is None."
+            )