AddExecutionViability transform (#4547)

Summary:

Transform that adds failure-awareness capability to Ax optimization.

This transform enables Ax to learn from deterministic trial failures (ABANDONED trials) and avoid sampling similar parameter configurations that are likely to fail. It achieves this by:

1. Adding a "execution_viable" metric to experiment data based on trial status
       - ABANDONED trials get feasibility value of 0.0 (not viable)
       - Other trials get feasibility value of 1.0 (viable)

2. Adding the execution_viable constraint to the optimization config
       - The constraint enforces P(execution_viable) >= threshold
       - This guides the acquisition function to avoid non-viable regions

Reviewed By: saitcakmak

Differential Revision: D85185246

Author: Sunny Shen

ax/adapter/transforms/add_execution_viability.py

@@ -0,0 +1,252 @@
+#!/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 __future__ import annotations
+
+from logging import Logger
+from typing import TYPE_CHECKING
+
+import pandas as pd
+from ax.adapter.data_utils import ExperimentData
+from ax.adapter.transforms.base import Transform
+from ax.core.base_trial import BaseTrial, TrialStatus
+from ax.core.metric import Metric
+from ax.core.observation import ObservationFeatures
+from ax.core.optimization_config import OptimizationConfig
+from ax.core.outcome_constraint import OutcomeConstraint
+from ax.core.types import ComparisonOp
+from ax.utils.common.constants import Keys
+from ax.utils.common.logger import get_logger
+
+if TYPE_CHECKING:
+    # import as module to make sphinx-autodoc-typehints happy
+    from ax import adapter as adapter_module  # noqa F401
+
+logger: Logger = get_logger(__name__)
+
+EXECUTION_VIABLE_METRIC_NAME = "execution_viable"
+
+
+class AddExecutionViability(Transform):
+    """Transform that adds failure-awareness capability to Ax optimization.
+
+    This transform enables Ax to learn from deterministic trial failures (ABANDONED
+    trials) and avoid sampling similar parameter configurations that are likely to
+    fail. It achieves this by:
+
+    1. Adding an "execution_viable" metric to experiment data based on trial status
+       - ABANDONED trials get execution_viable value of 0.0 (not viable)
+       - Other trials get execution_viable value of 1.0 (viable)
+
+    2. Adding an execution viability constraint to the optimization config
+       - The constraint enforces P(execution_viable) >= threshold
+       - This guides the acquisition function to avoid regions likely to fail
+
+    The transform only activates after observing a minimum number of ABANDONED trials
+    to ensure there is sufficient data to model the failure region. Before reaching
+    this threshold, the transform acts as a no-op.
+
+    Config options:
+        feasibility_threshold: float (default 0.8)
+            Minimum probability of execution viability required for new candidates.
+        min_abandoned_trials: int (default 3)
+            Minimum number of ABANDONED trials required before the transform activates.
+            If fewer than this many ABANDONED trials exist, the transform does nothing.
+
+    Example usage:
+        >>> transform = AddExecutionViability(
+        ...     config={
+        ...         "feasibility_threshold": 0.8,
+        ...         "min_abandoned_trials": 3,
+        ...     }
+        ... )
+        >>> # Transform adds execution viability constraint to optimization
+        >>> new_opt_config = transform.transform_optimization_config(opt_config)
+        >>> # Transform adds execution_viable metric to data
+        >>> transformed_data = transform.transform_experiment_data(exp_data)
+    """
+
+    @property
+    def min_abandoned_trials(self) -> int:
+        """Minimum ABANDONED trials required before the transform activates."""
+        raw_value = self.config.get("min_abandoned_trials", 3)
+        return int(raw_value) if isinstance(raw_value, (int, float)) else 3
+
+    def _should_activate(
+        self, adapter: adapter_module.base.Adapter
+    ) -> tuple[bool, int, list[BaseTrial]]:
+        """Check if transform should activate based on abandoned trial count.
+
+        Returns:
+            A tuple of (should_activate, abandoned_count, abandoned_trials)
+        """
+        experiment = adapter._experiment
+        abandoned_trials = experiment.trials_by_status.get(TrialStatus.ABANDONED, [])
+        abandoned_count = len(abandoned_trials)
+        should_activate = abandoned_count >= self.min_abandoned_trials
+        return should_activate, abandoned_count, abandoned_trials
+
+    def transform_experiment_data(
+        self, experiment_data: ExperimentData
+    ) -> ExperimentData:
+        """Transform experiment data to add execution viability metrics.
+
+        Only activates after observing at least min_abandoned_trials ABANDONED trials.
+        Returns the original data unchanged if this threshold is not met.
+
+        This method handles two types of ABANDONED trials:
+        1. ABANDONED trials WITH data: These already exist in
+           experiment_data and will get execution_viable = 0 added to their
+           existing observations.
+        2. ABANDONED trials WITHOUT data: These are missing from
+           experiment_data (e.g., trials that failed due to metric errors).
+           We add synthetic observations for these with execution_viable = 0 so
+           the model can learn about regions likely to fail.
+        """
+        if self.adapter is None:
+            raise ValueError(
+                "Adapter must be provided for AddExecutionViability transform."
+            )
+
+        adapter = self.adapter
+        should_activate, abandoned_count, abandoned_trials = self._should_activate(
+            adapter
+        )
+
+        if not should_activate:
+            logger.debug(
+                f"AddExecutionViability transform inactive: "
+                f"only {abandoned_count} ABANDONED trials observed "
+                f"(need {self.min_abandoned_trials}). Returning original data."
+            )
+            return experiment_data
+
+        # Proceed with adding execution viability metric
+        obs_data = experiment_data.observation_data
+        arm_data = experiment_data.arm_data
+
+        # Add rows for abandoned trials without data
+        trials_in_data = set(obs_data.index.get_level_values("trial_index").unique())
+        has_step = "step" in obs_data.index.names
+        new_arm_entries = []
+        for trial in abandoned_trials:
+            if trial.index not in trials_in_data:
+                for arm in trial.arms:
+                    idx = (
+                        (trial.index, arm.name, 1.0)
+                        if has_step
+                        else (trial.index, arm.name)
+                    )
+                    obs_data.loc[idx, :] = float("nan")
+
+                    arm_row_data = dict(arm.parameters)
+                    metadata_raw = trial._get_candidate_metadata(arm.name)
+                    metadata = metadata_raw if metadata_raw is not None else {}
+                    if (
+                        Keys.TRIAL_COMPLETION_TIMESTAMP not in metadata
+                        and trial._time_completed is not None
+                    ):
+                        metadata[Keys.TRIAL_COMPLETION_TIMESTAMP] = (
+                            trial._time_completed.timestamp()
+                        )
+                    arm_row_data["metadata"] = metadata  # pyre-ignore[6]
+                    new_arm_entries.append(
+                        {
+                            "trial_index": trial.index,
+                            "arm_name": arm.name,
+                            **arm_row_data,
+                        }
+                    )
+
+        if new_arm_entries:
+            new_arm_df = pd.DataFrame(new_arm_entries).set_index(
+                ["trial_index", "arm_name"]
+            )
+            arm_data = pd.concat([arm_data, new_arm_df])
+            logger.debug(
+                f"AddExecutionViability: Added synthetic observations for "
+                f"{len(new_arm_entries)} arms from ABANDONED trials without data"
+            )
+
+        # Assign viability for ALL rows in one shot
+        trial_indices = obs_data.index.get_level_values("trial_index")
+        abandoned_set = {t.index for t in abandoned_trials}
+        obs_data[("mean", EXECUTION_VIABLE_METRIC_NAME)] = [
+            0.0 if idx in abandoned_set else 1.0 for idx in trial_indices
+        ]
+        obs_data[("sem", EXECUTION_VIABLE_METRIC_NAME)] = float("nan")
+
+        logger.debug(
+            f"AddExecutionViability transform active: "
+            f"{abandoned_count} ABANDONED trials observed "
+            f"(threshold: {self.min_abandoned_trials})"
+        )
+
+        return ExperimentData(
+            arm_data=arm_data,
+            observation_data=obs_data,
+        )
+
+    def transform_optimization_config(
+        self,
+        optimization_config: OptimizationConfig,
+        adapter: adapter_module.base.Adapter | None = None,
+        fixed_features: ObservationFeatures | None = None,
+    ) -> OptimizationConfig:
+        """Transform optimization config to add execution viability constraint.
+
+        Only activates after observing at least min_abandoned_trials ABANDONED trials.
+        Returns the original config unchanged if this threshold is not met.
+        """
+        adapter = adapter or self.adapter
+        if adapter is None:
+            raise ValueError(
+                "Adapter must be provided for AddExecutionViability transform."
+            )
+
+        should_activate, abandoned_count, _ = self._should_activate(adapter)
+
+        if not should_activate:
+            logger.debug(
+                f"AddExecutionViability transform inactive: "
+                f"only {abandoned_count} ABANDONED trials observed "
+                f"(need {self.min_abandoned_trials}). Returning original config."
+            )
+            return optimization_config
+
+        # Proceed with adding execution viability constraint
+        viability_metric = Metric(
+            name=EXECUTION_VIABLE_METRIC_NAME,
+            lower_is_better=False,
+        )
+        viability_constraint = OutcomeConstraint(
+            metric=viability_metric,
+            op=ComparisonOp.GEQ,
+            bound=self.config.get("feasibility_threshold", 0.8),  # pyre-ignore [6]
+            relative=False,
+        )
+
+        # Create a new list with existing constraints plus the viability constraint
+        new_outcome_constraints = list(optimization_config.outcome_constraints)
+        new_outcome_constraints.append(viability_constraint)
+
+        transformed_opt_config = optimization_config.clone_with_args(
+            outcome_constraints=new_outcome_constraints,
+        )
+
+        # Add viability metric to outcomes if not already present
+        if viability_metric.name not in adapter.outcomes:
+            adapter.outcomes.append(viability_metric.name)
+
+        logger.debug(
+            f"AddExecutionViability constraint active: "
+            f"{abandoned_count} ABANDONED trials observed "
+            f"(threshold: {self.min_abandoned_trials})"
+        )
+
+        return transformed_opt_config