Clean up align_partial_results, eliminate double alignment, and remove Data round-trip

Summary:
Profiling-driven cleanup and performance improvements to the early stopping pipeline.

**1. Remove redundant validations from `align_partial_results` (utils.py)**

`align_partial_results` contained three validation/logging blocks that are redundant with upstream checks already performed by `_lookup_and_validate` in `base.py`:

- **Missing metrics check + `raise ValueError` (lines 166-168):** `_lookup_and_validate` already verifies each metric signature exists in the data (base.py lines 234-240) and returns `None` before `align_partial_results` is ever called.
- **Per-metric logging loop (lines 172-180):** The "no data" branch is unreachable because the upstream checks guarantee data exists for each metric. The debug logging about MAP_KEY ranges is misplaced for a pure data-alignment function.
- **Trial-to-arm uniqueness check (lines 186-198):** Redundant with `is_eligible_any` (base.py lines 344-352) which already rejects `BatchTrial` -- the only way a trial gets multiple arms.

The arm-to-trial uniqueness check was moved (not removed) to `_lookup_and_validate`, where it properly guards all downstream consumers, not just alignment.

The `isin` filter (`df = df[df['metric_signature'].isin(metrics)]`) was kept -- it is part of `align_partial_results`'s own contract since the function accepts a `metrics` argument and callers depend on it filtering to those metrics.

After cleanup, `align_partial_results` is a focused alignment function: filter to requested metrics -> drop `arm_name` -> drop duplicates -> sort -> pivot -> interpolate.

**2. Eliminate double data alignment in `PercentileEarlyStoppingStrategy` (percentile.py)**

When `check_safe=True`, the base class `should_stop_trials_early` calls `_is_harmful()` which calls `_prepare_aligned_frames()`, then `_should_stop_trials_early()` calls `_prepare_aligned_frames()` again -- identical data lookup + alignment running twice.

Fix: override `should_stop_trials_early` to call `_default_objective_and_direction()` and `_prepare_aligned_frames()` once, passing results as explicit keyword arguments (`metric_signature`, `minimize`, `aligned_frames`) to both `_is_harmful` and `_should_stop_trials_early`. Each defaults to `None` and falls back to computing from scratch when not provided, preserving backward compatibility for subclasses.

**3. Eliminate the `Data` round-trip in `_lookup_and_validate` (base.py, threshold.py)**

`_lookup_and_validate` (formerly `_lookup_and_validate_data`) used to wrap its filtered DataFrame back into `Data(df=filtered_df)` at the end, purely to satisfy its `Data | None` return type. The sole consumer (`_prepare_aligned_frames`) immediately called `.full_df` to get the DataFrame back. This triggered a pointless df-to-DataRow-to-df round-trip on every call:

```
DataFrame -> itertuples -> list[DataRow] -> from_records -> regex sort -> cast -> DataFrame
```

Profiling shows this round-trip is **100% overhead** at every scale:

| Scale           | Rows    | Round-trip cost | Replacement (df.copy) |
|-----------------|---------|----------------|-----------------------|
| tiny (5x10)     | 50      | 4ms            | <0.1ms                |
| typical (20x100)| 2,000   | 15ms           | <0.1ms                |
| large (50x200)  | 10,000  | 60ms           | <0.1ms                |
| xlarge (100x200)| 20,000  | 160ms          | <0.1ms                |
| huge (200x500)  | 100,000 | 733ms          | <0.1ms                |

The cost comes from `Data.__init__` iterating every row via `itertuples()` to build `list[DataRow]` (~35ms at 10k rows), then `Data.full_df` reconstructing the DataFrame via `from_records()` (~12ms) and running regex-based arm name parsing in `sort_by_trial_index_and_arm_name()` (~19ms) -- none of which is needed since we already had the DataFrame.

Fix: change `_lookup_and_validate` to return `pd.DataFrame | None` directly. Update `ModelBasedEarlyStoppingStrategy` and `ThresholdEarlyStoppingStrategy` accordingly.

End-to-end profiling of `should_stop_trials_early` at the 50x200 scale (10k rows) shows a ~19% speedup (323ms -> ~263ms), with the benefit growing at larger scales (up to ~41% at 100k rows).

**4. Naming cleanup**

Renamed methods and variables to disambiguate between `Data` objects and `pd.DataFrame`:
- `_lookup_and_validate_data` -> `_lookup_and_validate` (returns `pd.DataFrame | None`)
- `_prepare_aligned_data` -> `_prepare_aligned_frames` (returns a tuple of DataFrames)
- Variables holding `_lookup_and_validate` results: `data`/`map_data`/`data_lookup` -> `df`
- Parameter `aligned_data` -> `aligned_frames`

**5. Tests**

- Removed tests for moved/removed `align_partial_results` validations.
- Added test for the arm-to-trial uniqueness check in `_lookup_and_validate`.
- Added profiling notebook at `ax/early_stopping/profiling.ipynb`.

Differential Revision: D98544835

Author: ltiao

ax/early_stopping/strategies/base.py

@@ -15,7 +15,7 @@
 import pandas as pd
 from ax.adapter.data_utils import _maybe_normalize_map_key
 from ax.core.batch_trial import BatchTrial
-from ax.core.data import Data, MAP_KEY
+from ax.core.data import MAP_KEY
 from ax.core.experiment import Experiment
 from ax.core.trial_status import TrialStatus
 from ax.early_stopping.utils import (
@@ -217,13 +217,10 @@ def estimate_early_stopping_savings(self, experiment: Experiment) -> float:
 
         return estimate_early_stopping_savings(experiment=experiment)
 
-    def _lookup_and_validate_data(
+    def _lookup_and_validate(
         self, experiment: Experiment, metric_signatures: list[str]
-    ) -> Data | None:
-        """Looks up and validates the `Data` used for early stopping that
-        is associated with `metric_signatures`. This function also handles normalizing
-        progressions.
-        """
+    ) -> pd.DataFrame | None:
+        """Look up and validate experiment data for early stopping."""
         data = experiment.lookup_data()
         if data.df.empty:
             logger.info(
@@ -250,6 +247,17 @@ def _lookup_and_validate_data(
         full_df = data.full_df
         full_df = full_df[full_df["metric_signature"].isin(metric_signatures)]
 
+        # Check that no arm name appears across multiple trials.
+        # This can happen with duplicate arm parameterizations that reuse arm
+        # names across trials, which would corrupt the alignment step.
+        arm_trial_counts = full_df.groupby("arm_name")["trial_index"].nunique()
+        bad_arms = arm_trial_counts[arm_trial_counts > 1]
+        if len(bad_arms) > 0:
+            raise UnsupportedError(
+                f"Arm(s) {bad_arms.index.tolist()} appear across multiple "
+                f"trial indices. Each arm name must map to exactly one trial."
+            )
+
         # Drop rows with NaN values in MAP_KEY column to prevent issues in
         # align_partial_results which uses MAP_KEY as the pivot index
         nan_mask = full_df[MAP_KEY].isna()
@@ -264,7 +272,7 @@ def _lookup_and_validate_data(
 
         if self.normalize_progressions:
             full_df = _maybe_normalize_map_key(df=full_df)
-        return Data(df=full_df)
+        return full_df
 
     @staticmethod
     def _log_and_return_no_data(
@@ -547,7 +555,7 @@ def _all_objectives_and_directions(self, experiment: Experiment) -> dict[str, bo
 
         return directions
 
-    def _prepare_aligned_data(
+    def _prepare_aligned_frames(
         self, experiment: Experiment, metric_signatures: list[str]
     ) -> tuple[pd.DataFrame, pd.DataFrame] | None:
         """Get raw experiment data and align it for early stopping evaluation.
@@ -564,15 +572,15 @@ def _prepare_aligned_data(
               with first level ["mean", "sem"] and second level metric signatures
             Returns None if data cannot be retrieved or aligned.
         """
-        data = self._lookup_and_validate_data(
+        long_df = self._lookup_and_validate(
             experiment=experiment, metric_signatures=metric_signatures
         )
-        if data is None:
+        if long_df is None:
             return None
 
         try:
             multilevel_wide_df = align_partial_results(
-                df=(long_df := data.full_df),
+                df=long_df,
                 metrics=metric_signatures,
             )
         except Exception as e:
@@ -651,18 +659,15 @@ def __init__(
         )
         self.min_progression_modeling = min_progression_modeling
 
-    def _lookup_and_validate_data(
+    def _lookup_and_validate(
         self, experiment: Experiment, metric_signatures: list[str]
-    ) -> Data | None:
-        """Looks up and validates the `Data` used for early stopping that
-        is associated with `metric_signatures`. This function also handles normalizing
-        progressions.
+    ) -> pd.DataFrame | None:
+        """Look up and validate experiment data for early stopping, applying
+        ``min_progression_modeling`` filter if configured.
         """
-        map_data = super()._lookup_and_validate_data(
+        df = super()._lookup_and_validate(
             experiment=experiment, metric_signatures=metric_signatures
         )
-        if map_data is not None and self.min_progression_modeling is not None:
-            full_df = map_data.full_df
-            full_df = full_df[full_df[MAP_KEY] >= self.min_progression_modeling]
-            map_data = Data(df=full_df)
-        return map_data
+        if df is not None and self.min_progression_modeling is not None:
+            df = df[df[MAP_KEY] >= self.min_progression_modeling]
+        return df

ax/early_stopping/strategies/percentile.py

@@ -124,10 +124,63 @@ def __init__(
                 "with multiple metrics."
             )
 
+    def should_stop_trials_early(
+        self,
+        trial_indices: set[int],
+        experiment: Experiment,
+        current_node: GenerationNode | None = None,
+    ) -> dict[int, str | None]:
+        """Decide whether trials should be stopped before evaluation is fully concluded.
+
+        Overrides the base class to compute aligned data once and reuse it for
+        both the safety check (``_is_harmful``) and the stopping decision
+        (``_should_stop_trials_early``), avoiding redundant data lookups and
+        alignment when ``check_safe=True``.
+
+        Args:
+            trial_indices: Indices of candidate trials to stop early.
+            experiment: Experiment that contains the trials and other contextual data.
+            current_node: The current ``GenerationNode`` on the ``GenerationStrategy``
+                used to generate trials for the ``Experiment``.
+
+        Returns:
+            A dictionary mapping trial indices that should be early stopped to
+            (optional) messages with the associated reason.
+        """
+        metric_signature, minimize = self._default_objective_and_direction(
+            experiment=experiment
+        )
+        aligned_frames = self._prepare_aligned_frames(
+            experiment=experiment, metric_signatures=[metric_signature]
+        )
+        if aligned_frames is None:
+            return {}
+
+        if self.check_safe and self._is_harmful(
+            trial_indices=trial_indices,
+            experiment=experiment,
+            metric_signature=metric_signature,
+            minimize=minimize,
+            aligned_frames=aligned_frames,
+        ):
+            return {}
+
+        return self._should_stop_trials_early(
+            trial_indices=trial_indices,
+            experiment=experiment,
+            current_node=current_node,
+            metric_signature=metric_signature,
+            minimize=minimize,
+            aligned_frames=aligned_frames,
+        )
+
     def _is_harmful(
         self,
         trial_indices: set[int],
         experiment: Experiment,
+        metric_signature: str | None = None,
+        minimize: bool | None = None,
+        aligned_frames: tuple[pd.DataFrame, pd.DataFrame] | None = None,
     ) -> bool:
         """Check if the early stopping strategy would stop the globally best trial.
 
@@ -139,21 +192,30 @@ def _is_harmful(
         Args:
             trial_indices: Set of trial indices being evaluated (ignored).
             experiment: Experiment that contains the trials and other contextual data.
+            metric_signature: The metric signature to evaluate. If not provided,
+                it is inferred from the experiment's objective.
+            minimize: Whether the metric is being minimized. If not provided,
+                it is inferred from the experiment's objective.
+            aligned_frames: Pre-computed ``(long_df, multilevel_wide_df)`` tuple
+                from ``_prepare_aligned_frames``. If not provided, it is computed
+                from scratch.
 
         Returns:
             True if the strategy would have stopped the globally best trial,
             False otherwise.
         """
-        metric_signature, minimize = self._default_objective_and_direction(
-            experiment=experiment
-        )
-        maybe_aligned_dataframes = self._prepare_aligned_data(
-            experiment=experiment, metric_signatures=[metric_signature]
-        )
-        if maybe_aligned_dataframes is None:
-            return False
+        if metric_signature is None or minimize is None:
+            metric_signature, minimize = self._default_objective_and_direction(
+                experiment=experiment
+            )
+        if aligned_frames is None:
+            aligned_frames = self._prepare_aligned_frames(
+                experiment=experiment, metric_signatures=[metric_signature]
+            )
+            if aligned_frames is None:
+                return False
 
-        long_df, multilevel_wide_df = maybe_aligned_dataframes
+        long_df, multilevel_wide_df = aligned_frames
         wide_df = multilevel_wide_df["mean"][metric_signature]
 
         # Get completed trials
@@ -180,6 +242,9 @@ def _should_stop_trials_early(
         trial_indices: set[int],
         experiment: Experiment,
         current_node: GenerationNode | None = None,
+        metric_signature: str | None = None,
+        minimize: bool | None = None,
+        aligned_frames: tuple[pd.DataFrame, pd.DataFrame] | None = None,
     ) -> dict[int, str | None]:
         """Stop a trial if its performance is in the bottom `percentile_threshold`
         of the trials at the same step.
@@ -191,22 +256,31 @@ def _should_stop_trials_early(
                 used to generate trials for the ``Experiment``. Early stopping
                 strategies may utilize components of the current node when making
                 stopping decisions.
+            metric_signature: The metric signature to evaluate. If not provided,
+                it is inferred from the experiment's objective.
+            minimize: Whether the metric is being minimized. If not provided,
+                it is inferred from the experiment's objective.
+            aligned_frames: Pre-computed ``(long_df, multilevel_wide_df)`` tuple
+                from ``_prepare_aligned_frames``. If not provided, it is computed
+                from scratch.
 
         Returns:
             A dictionary mapping trial indices that should be early stopped to
             (optional) messages with the associated reason. An empty dictionary
             means no suggested updates to any trial's status.
         """
-        metric_signature, minimize = self._default_objective_and_direction(
-            experiment=experiment
-        )
-        maybe_aligned_dataframes = self._prepare_aligned_data(
-            experiment=experiment, metric_signatures=[metric_signature]
-        )
-        if maybe_aligned_dataframes is None:
-            return {}
+        if metric_signature is None or minimize is None:
+            metric_signature, minimize = self._default_objective_and_direction(
+                experiment=experiment
+            )
+        if aligned_frames is None:
+            aligned_frames = self._prepare_aligned_frames(
+                experiment=experiment, metric_signatures=[metric_signature]
+            )
+            if aligned_frames is None:
+                return {}
 
-        long_df, multilevel_wide_df = maybe_aligned_dataframes
+        long_df, multilevel_wide_df = aligned_frames
         wide_df = multilevel_wide_df["mean"][metric_signature]
 
         # default checks on `min_progression` and `min_curves`; if not met, don't do

ax/early_stopping/strategies/threshold.py

@@ -112,15 +112,13 @@ def _should_stop_trials_early(
         metric_signature, minimize = self._default_objective_and_direction(
             experiment=experiment
         )
-        data = self._lookup_and_validate_data(
+        df = self._lookup_and_validate(
             experiment=experiment, metric_signatures=[metric_signature]
         )
-        if data is None:
+        if df is None:
             # don't stop any trials if we don't get data back
             return {}
 
-        df = data.full_df
-
         # default checks on `min_progression` and `min_curves`; if not met, don't do
         # early stopping at all and return {}
         if not self.is_eligible_any(