Extract MapDataReplayState for multi-metric experiment replay (#5084)

Summary:

The experiment replay system (`MapDataReplayMetric`, `MapDataReplayRunner`, `replay_experiment`) is hardcoded for single-objective optimization, blocking multi-objective early stopping. `MapDataReplayMetric` conflates data serving with progression state (offset, scaling factor, per-trial step counters), so multiple metrics cannot share a coherent timeline. The scaling factor -- `mean((last_step - offset) / num_points)` -- is an unmotivated heuristic that distorts cross-trial timing. The runner is tightly coupled to a single metric instance.

This diff extracts shared state into a `MapDataReplayState` coordinator. The runner owns progression, the metric is a read-only accessor.

**Key changes:**

- **`MapDataReplayState`** (new): Normalized cursor model. Computes global `min_prog`/`max_prog` across all metrics and trials, advances per-trial cursors by fixed `step_size`, maps to raw progression via `curr_prog = min_prog + cursor * (max_prog - min_prog)`. Serves original MAP_KEY values; downstream ESS normalizes independently.
- **`MapDataReplayMetric`** (simplified): Thin wrapper holding a state reference and `metric_signature`. Delegates `fetch_trial_data` to `state.get_data()`. All state-owning attributes removed.
- **`MapDataReplayRunner`** (simplified): Takes shared state instead of metric. Calls `advance_trial()` for running trials.
- **`replay_experiment`**: Accepts `metrics: list[Metric]`. Builds `OptimizationConfig` or `MultiObjectiveOptimizationConfig` accordingly. Extracts objective thresholds from the historical experiment's config. Re-indexes non-contiguous trial indices to contiguous `0, 1, 2, ...`. Deprecates `num_samples_per_curve` (superseded by `step_size`).
- **`estimate_hypothetical_early_stopping_savings`**: Updated for `metrics: list[Metric]`.
- **Removed**: `_compute_trial_stats`, `_compute_scaling_factor` (superseded by cursor model).
- **Downstream consumers updated**: `ax_sweep_orchestrator`, `early_stopping_healthcheck`.

Differential Revision: D96999702

Author: ltiao

ax/analysis/healthcheck/early_stopping_healthcheck.py

@@ -413,7 +413,7 @@ def _report_early_stopping_nudge(
         try:
             savings = estimate_hypothetical_early_stopping_savings(
                 experiment=experiment,
-                metric=metric,
+                metrics=[metric],
                 max_pending_trials=self.max_pending_trials,
             )
         except Exception as e:

ax/early_stopping/experiment_replay.py

@@ -7,14 +7,20 @@
 # pyre-strict
 
 import logging
+import warnings
 from logging import Logger
 from time import perf_counter
 
 from ax.adapter.registry import Generators
+from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.metric import Metric
-from ax.core.objective import Objective
-from ax.core.optimization_config import OptimizationConfig
+from ax.core.objective import MultiObjective, Objective
+from ax.core.optimization_config import (
+    MultiObjectiveOptimizationConfig,
+    OptimizationConfig,
+)
+from ax.core.outcome_constraint import OutcomeConstraint
 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
@@ -25,7 +31,7 @@
     GenerationStep,
     GenerationStrategy,
 )
-from ax.metrics.map_replay import MapDataReplayMetric
+from ax.metrics.map_replay import MapDataReplayMetric, MapDataReplayState
 from ax.orchestration.orchestrator import Orchestrator, OrchestratorOptions
 from ax.runners.map_replay import MapDataReplayRunner
 from ax.utils.common.logger import get_logger
@@ -43,16 +49,36 @@ def replay_experiment(
     historical_experiment: Experiment,
     num_samples_per_curve: int,
     max_replay_trials: int,
-    metric: Metric,
+    metrics: list[Metric],
     max_pending_trials: int,
     early_stopping_strategy: BaseEarlyStoppingStrategy | None,
     logging_level: int = logging.ERROR,
 ) -> Experiment | None:
-    """A utility function for replaying a historical experiment's data
-    by initializing a Orchestrator that quickly steps through the existing data.
-    The main purpose of this function is to compute an hypothetical capacity
-    savings for a given `early_stopping_strategy`.
+    """Replay a historical experiment's data through an Orchestrator.
+
+    Initializes an Orchestrator that steps through existing data to compute
+    hypothetical capacity savings for a given ``early_stopping_strategy``.
+    Supports both single-objective and multi-objective optimization.
+
+    Args:
+        historical_experiment: The experiment whose data to replay.
+        num_samples_per_curve: Deprecated. Number of samples per curve for
+            subsampling. Use ``step_size`` on ``MapDataReplayState`` instead.
+        max_replay_trials: Maximum number of trials to replay.
+        metrics: List of metrics to replay. For single-objective, provide
+            one metric. For multi-objective, provide multiple metrics.
+        max_pending_trials: Maximum number of pending trials for the
+            replay orchestrator.
+        early_stopping_strategy: The early stopping strategy to evaluate.
+        logging_level: Logging level for the orchestrator.
     """
+    warnings.warn(
+        "The `num_samples_per_curve` parameter is deprecated and will be "
+        "removed in a future release. The `step_size` parameter on "
+        "`MapDataReplayState` controls replay granularity.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     historical_map_data = historical_experiment.lookup_data()
     if not historical_map_data.has_step_column:
         logger.warning(
@@ -62,16 +88,51 @@ def replay_experiment(
     historical_map_data = historical_map_data.subsample(
         limit_rows_per_group=num_samples_per_curve, include_first_last=True
     )
-    replay_metric = MapDataReplayMetric(
-        name=f"replay_{historical_experiment.name}",
-        map_data=historical_map_data,
-        metric_name=metric.name,
-        lower_is_better=metric.lower_is_better,
-    )
-    optimization_config = OptimizationConfig(
-        objective=Objective(metric=replay_metric),
+
+    # Re-index non-contiguous trial indices to contiguous 0, 1, 2, ...
+    # so that replay trial N maps to the Nth historical trial.
+    df = historical_map_data.full_df
+    sorted_trial_indices = sorted(df["trial_index"].unique())
+    trial_index_map = {old: new for new, old in enumerate(sorted_trial_indices)}
+    df = df.copy()
+    df["trial_index"] = df["trial_index"].map(trial_index_map)
+    historical_map_data = Data(df=df)
+
+    metric_signatures = [m.signature for m in metrics]
+    replay_state = MapDataReplayState(
+        map_data=historical_map_data, metric_signatures=metric_signatures
     )
-    runner = MapDataReplayRunner(replay_metric=replay_metric)
+
+    replay_metrics = [
+        MapDataReplayMetric(
+            name=m.name,
+            replay_state=replay_state,
+            metric_signature=m.signature,
+            lower_is_better=m.lower_is_better,
+        )
+        for m in metrics
+    ]
+
+    if len(replay_metrics) == 1:
+        optimization_config: OptimizationConfig = OptimizationConfig(
+            objective=Objective(metric=replay_metrics[0]),
+        )
+    else:
+        # Extract objective thresholds from the historical experiment's config
+        historical_opt_config = historical_experiment.optimization_config
+        objective_thresholds: list[OutcomeConstraint] = []
+        if isinstance(historical_opt_config, MultiObjectiveOptimizationConfig):
+            objective_thresholds = [
+                ot.clone() for ot in historical_opt_config.objective_thresholds
+            ]
+        optimization_config = MultiObjectiveOptimizationConfig(
+            objective=MultiObjective(
+                objectives=[Objective(metric=m) for m in replay_metrics]
+            ),
+            objective_thresholds=objective_thresholds,
+        )
+
+    runner = MapDataReplayRunner(replay_state=replay_state)
 
     # Setup a new experiment with a dummy search space
     dummy_search_space = SearchSpace(
@@ -89,10 +150,10 @@ def replay_experiment(
         optimization_config=optimization_config,
         search_space=dummy_search_space,
         runner=runner,
-        metrics=[replay_metric],
+        metrics=replay_metrics,
     )
 
-    # Setup a Orchestrator with a dummy gs to replay the historical experiment
+    # Setup an Orchestrator with a dummy gs to replay the historical experiment
     dummy_sobol_gs = GenerationStrategy(
         name="sobol",
         steps=[
@@ -101,7 +162,7 @@ def replay_experiment(
     )
     options = OrchestratorOptions(
         max_pending_trials=max_pending_trials,
-        total_trials=min(len(historical_experiment.trials), max_replay_trials),
+        total_trials=min(len(sorted_trial_indices), max_replay_trials),
         seconds_between_polls_backoff_factor=1.0,
         min_seconds_before_poll=0.0,
         init_seconds_between_polls=0,
@@ -119,7 +180,7 @@ def replay_experiment(
 
 def estimate_hypothetical_early_stopping_savings(
     experiment: Experiment,
-    metric: Metric,
+    metrics: list[Metric],
     max_pending_trials: int = MAX_PENDING_TRIALS,
 ) -> float:
     """Estimate hypothetical early stopping savings using experiment replay.
@@ -130,7 +191,7 @@ def estimate_hypothetical_early_stopping_savings(
 
     Args:
         experiment: The experiment to analyze.
-        metric: The metric to use for early stopping replay.
+        metrics: The metrics to use for early stopping replay.
         max_pending_trials: Maximum number of pending trials for the replay
             orchestrator. Defaults to 5.
 
@@ -156,7 +217,7 @@ def estimate_hypothetical_early_stopping_savings(
         historical_experiment=experiment,
         num_samples_per_curve=REPLAY_NUM_POINTS_PER_CURVE,
         max_replay_trials=MAX_REPLAY_TRIALS,
-        metric=metric,
+        metrics=metrics,
         max_pending_trials=max_pending_trials,
         early_stopping_strategy=default_ess,
     )