Decouple metric fetch errors from trial status in Orchestrator (#5287)

Summary:
Pull Request resolved: #5287

Keep trial execution status authoritative to runner polling when metric fetching fails. Orchestrator now leaves completed trials completed, uses MetricAvailability to count incomplete optimization-config data in the existing failure-rate policy, and handles expanded multi-metric result keys without assuming they are registered experiment metrics.

This mirrors the behavior core Ax already has in the Orchestrator.

This implements the core design from the abandoned RFC D98741656 and fixes T287759832.

Reviewed By: ltiao, esantorella

Differential Revision: D119207750

fbshipit-source-id: 040c93063541bd15388a64c7fef649fcc7d6c860

Author: saitcakmak

ax/core/metric.py

@@ -88,10 +88,8 @@ class Metric(SortableBase, SerializationMixin):
         properties: Properties specific to a particular metric.
     """
 
-    # The set of exception types stored in a ``MetchFetchE.exception`` that are
-    # recoverable ``orchestrator._fetch_and_process_trials_data_results()``.
-    # Exception may be a subclass of any of these types.  If you want your metric
-    # to never fail the trial, set this to ``{Exception}`` in your metric subclass.
+    # Legacy configuration retained for compatibility with existing Metric subclasses.
+    # Metric fetching errors no longer change trial status in Orchestrator.
     recoverable_exceptions: set[type[Exception]] = set()
     has_map_data: bool = False
 
@@ -166,9 +164,7 @@ def period_of_new_data_after_trial_completion(cls) -> timedelta:
 
     @classmethod
     def is_recoverable_fetch_e(cls, metric_fetch_e: MetricFetchE) -> bool:
-        """Checks whether the given MetricFetchE is recoverable for this metric class
-        in ``orchestrator._fetch_and_process_trials_data_results``.
-        """
+        """Check whether the given MetricFetchE is recoverable for this metric."""
         if metric_fetch_e.exception is None:
             return False
         return any(

ax/orchestration/orchestrator.py

@@ -32,6 +32,7 @@
 from ax.core.runner import Runner
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
+from ax.core.utils import compute_metric_availability, MetricAvailability
 from ax.exceptions.core import (
     AxError,
     DataRequiredError,
@@ -71,20 +72,20 @@
 trials attached to the underlying Ax experiment '{experiment_name}'.
 """
 FAILURE_EXCEEDED_MSG = (
-    "NOTE: This error is usually not caused by Ax. Please please check any trial "
-    "evaluation processes/jobs to see why they are failing, and ensure that they "
-    "succeed over the entire range of the parameters defined in this optimization.\n\n"
-    "Trials are failing or being abandoned at a rate {observed_rate} that exceeds the "
-    "tolerated trial failure rate of {f_rate} (at least {n_failed} out of first "
-    "{n_ran} trials failed or were abandoned). Checks are triggered both at the end "
-    "of an optimization and if at least {min_failed} trials have been "
-    "failed/abandoned, potentially automatically due to issues with the trial."
+    "NOTE: This error is usually not caused by Ax. Please check any trial evaluation "
+    "processes/jobs and metric fetching infrastructure to determine why trials are "
+    "failing or required data is missing.\n\n"
+    "Trials are FAILED or ABANDONED, or have incomplete data for metrics in the "
+    "optimization config, at a rate {observed_rate} that exceeds the tolerated trial "
+    "failure rate of {f_rate} "
+    "(at least {n_failed} out of the first {n_ran} trials). Checks are triggered both "
+    "at the end of an optimization and after at least {min_failed} affected trials."
 )
 METRIC_FETCH_ERR_MESSAGE = (
-    "A majority of the trial failures encountered are due to metric fetching errors. "
-    "This could mean the metrics are flaky, broken, or misconfigured. Please check "
-    "that the trial processes/jobs are successfully producing the expected metrics and "
-    "that the metric is correctly configured."
+    "One or more completed trials have incomplete data for metrics in the optimization "
+    "config. This could mean the metrics are flaky, broken, or misconfigured. Please "
+    "check that the trial processes/jobs are producing the expected metrics and that "
+    "the metrics are correctly configured."
 )
 
 EXPECTED_STAGED_MSG = (
@@ -191,17 +192,20 @@ class Orchestrator(WithDBSettingsBase, BestPointMixin):
     # Saved as a property so that it can be accessed after optimization is complex (ex.
     # for global stopping saving calculation).
     _num_remaining_requested_trials: int = 0
-    # Total number of MetricFetchEs encountered during the course of optimization. Note
-    # this is different from and may be greater than the number of trials that have
-    # been marked either FAILED or ABANDONED due to metric fetching errors.
+    # Total number of MetricFetchEs encountered for non-running trials during the
+    # course of optimization.
     _num_metric_fetch_e_encountered: int = 0
-    # Number of trials that have been marked either FAILED or ABANDONED due to
-    # MetricFetchE being encountered during _fetch_and_process_trials_data_results
+    # Number of completed trials with incomplete data for metrics in the optimization
+    # config.
+    # Retained under its existing name for telemetry compatibility.
     _num_trials_bad_due_to_err: int = 0
     # Keeps track of whether the allowed failure rate has been exceeded during
     # the optimization. If true, allows any pending trials to finish and raises
     # an error through self._complete_optimization.
     _failure_rate_has_been_exceeded: bool = False
+    # Counts captured when the failure rate first exceeds the tolerance. Pending
+    # trials may finish afterward, so the final error must use this snapshot.
+    _failure_rate_exceeded_counts: tuple[int, int] | None = None
     # Timestamp of last optimization start time (milliseconds since Unix epoch);
     # recorded in each `run_n_trials`.
     _latest_optimization_start_timestamp: int | None = None
@@ -1079,14 +1083,15 @@ def _check_if_failure_rate_exceeded(self, force_check: bool = False) -> bool:
         """Checks if the failure rate (set in Orchestrator options) has been exceeded at
         any point during the optimization.
 
-        NOTE: Both FAILED and ABANDONED trial statuses count towards the failure rate.
+        FAILED and ABANDONED statuses count towards the failure rate, as do completed
+        trials with incomplete data for metrics in the optimization config.
 
         Args:
             force_check: Indicates whether to force a failure-rate check
                 regardless of the number of trials that have been executed. If False
-                (default), the check will be skipped if the optimization has fewer than
-                five failed trials. If True, the check will be performed unless there
-                are 0 failures.
+                (default), the check will be skipped if the optimization has fewer
+                affected trials than ``min_failed_trials_for_failure_rate_check``. If
+                True, the check will be performed unless there are 0 affected trials.
 
         Effect on state:
             If the failure rate has been exceeded, a warning is logged and the private
@@ -1100,7 +1105,11 @@ def _check_if_failure_rate_exceeded(self, force_check: bool = False) -> bool:
         if self._failure_rate_has_been_exceeded:
             return True
 
-        num_bad_in_orchestrator = self._num_bad_in_orchestrator()
+        num_failed_or_abandoned = self._num_bad_in_orchestrator()
+        self._num_trials_bad_due_to_err = self._num_metric_incomplete_in_orchestrator()
+        num_bad_in_orchestrator = (
+            num_failed_or_abandoned + self._num_trials_bad_due_to_err
+        )
         # skip check if 0 failures
         if num_bad_in_orchestrator == 0:
             return False
@@ -1120,15 +1129,18 @@ def _check_if_failure_rate_exceeded(self, force_check: bool = False) -> bool:
         ) > self.options.tolerated_trial_failure_rate
 
         if failure_rate_exceeded:
-            if self._num_trials_bad_due_to_err > num_bad_in_orchestrator / 2:
+            if self._num_trials_bad_due_to_err > 0:
                 self.logger.warning(
-                    "MetricFetchE INFO: Sweep aborted due to an exceeded error rate, "
-                    "which was primarily caused by failure to fetch metrics. Please "
-                    "check if anything could cause your metrics to be flaky or "
-                    "broken."
+                    "MetricFetchE INFO: Sweep aborted due to an exceeded error rate "
+                    "that includes incomplete data for metrics in the optimization "
+                    "config."
                 )
             # NOTE: this private attribute causes `_get_max_pending_trials` to
             # return zero, which causes no further trials to be scheduled.
+            self._failure_rate_exceeded_counts = (
+                num_bad_in_orchestrator,
+                num_ran_in_orchestrator,
+            )
             self._failure_rate_has_been_exceeded = True
             return True
 
@@ -1138,19 +1150,23 @@ def error_if_failure_rate_exceeded(self, force_check: bool = False) -> None:
         """Raises an exception if the failure rate (set in Orchestrator options) has
         been exceeded at any point during the optimization.
 
-        NOTE: Both FAILED and ABANDONED trial statuses count towards the failure rate.
+        FAILED and ABANDONED statuses count towards the failure rate, as do completed
+        trials with incomplete data for metrics in the optimization config.
 
         Args:
             force_check: Indicates whether to force a failure-rate check
                 regardless of the number of trials that have been executed. If False
-                (default), the check will be skipped if the optimization has fewer than
-                five failed trials. If True, the check will be performed unless there
-                are 0 failures.
+                (default), the check will be skipped if the optimization has fewer
+                affected trials than ``min_failed_trials_for_failure_rate_check``. If
+                True, the check will be performed unless there are 0 affected trials.
         """
         if self._check_if_failure_rate_exceeded(force_check=force_check):
+            num_bad_in_orchestrator, num_ran_in_orchestrator = none_throws(
+                self._failure_rate_exceeded_counts
+            )
             raise self._get_failure_rate_exceeded_error(
-                num_bad_in_orchestrator=self._num_bad_in_orchestrator(),
-                num_ran_in_orchestrator=self._num_ran_in_orchestrator(),
+                num_bad_in_orchestrator=num_bad_in_orchestrator,
+                num_ran_in_orchestrator=num_ran_in_orchestrator,
             )
 
     def _error_if_status_quo_infeasible(self) -> None:
@@ -1403,8 +1419,8 @@ def _fetch_data_and_return_trial_indices_with_new_data(
             )
             return {
                 i
-                for i, results_by_metric_name in results.items()
-                for r in results_by_metric_name.values()
+                for i, results_by_metric_signature in results.items()
+                for r in results_by_metric_signature.values()
                 if r.is_ok()
             }
         return set()
@@ -1425,6 +1441,26 @@ def _num_ran_in_orchestrator(self) -> int:
         """Returns the number of trials that have been run by the orchestrator."""
         return len(self.experiment.trials) - self._num_preexisting_trials
 
+    def _num_metric_incomplete_in_orchestrator(self) -> int:
+        """Count completed trials with incomplete optimization config data."""
+        if self.experiment.optimization_config is None:
+            return 0
+
+        completed_trial_indices = [
+            trial.index
+            for trial in self.trials
+            if trial.status == TrialStatus.COMPLETED
+            and trial.index >= self._num_preexisting_trials
+        ]
+        metric_availability = compute_metric_availability(
+            experiment=self.experiment,
+            trial_indices=completed_trial_indices,
+        )
+        return sum(
+            availability != MetricAvailability.COMPLETE
+            for availability in metric_availability.values()
+        )
+
     def _apply_trial_statuses(
         self, polled_status_to_trial_idcs: dict[TrialStatus, set[int]]
     ) -> set[int]:
@@ -2071,8 +2107,7 @@ def _fetch_and_process_trials_data_results(
         trial_indices: Iterable[int],
     ) -> dict[int, dict[str, MetricFetchResult]]:
         """
-        Fetches results from experiment and modifies trial statuses depending on
-        success or failure.
+        Fetch results and report errors without changing trial statuses.
         """
 
         try:
@@ -2092,8 +2127,8 @@ def _fetch_and_process_trials_data_results(
             )
             return {}
 
-        for trial_index, results_by_metric_name in results.items():
-            for metric_name, result in results_by_metric_name.items():
+        for trial_index, results_by_metric_signature in results.items():
+            for metric_signature, result in results_by_metric_signature.items():
                 # If the fetch call succeeded, continue.
                 if result.is_ok():
                     continue
@@ -2102,59 +2137,34 @@ def _fetch_and_process_trials_data_results(
                 # we do not do anything
                 metric_fetch_e = result.unwrap_err()
 
-                # If the metric is available while running just continue (we can try
-                # again later).
-                # NOTE: We don't need to report fetching errors in this case either
-                metric = self.experiment.metrics[metric_name]
                 status = self.experiment.trials[trial_index].status
+                metric = self.experiment.signature_to_metric.get(metric_signature)
                 if (
-                    metric.is_available_while_running()
-                    and status == TrialStatus.RUNNING
+                    status == TrialStatus.RUNNING
+                    and metric is not None
+                    and metric.is_available_while_running()
                 ):
                     self.logger.info(
-                        f"MetricFetchE INFO: Because {metric_name} is "
-                        f"available_while_running and trial {trial_index} is still "
-                        "RUNNING continuing the experiment and retrying on next "
-                        "poll..."
+                        f"MetricFetchE INFO: Because trial {trial_index} is still "
+                        "RUNNING, continuing after failing to fetch "
+                        f"{metric_signature} and retrying on the next poll."
                     )
                     continue
 
                 self.logger.error(
-                    f"Failed to fetch {metric_name} for trial {trial_index} with "
+                    f"Failed to fetch {metric_signature} for trial {trial_index} with "
                     f"status {status}, found {metric_fetch_e}."
                 )
                 self._num_metric_fetch_e_encountered += 1
                 self._report_metric_fetch_e(
                     trial=self.experiment.trials[trial_index],
-                    metric_name=metric_name,
+                    metric_signature=metric_signature,
                     metric_fetch_e=metric_fetch_e,
                 )
 
-                # If the fetch failure was for a metric in the optimization config (an
-                # objective or constraint) mark the trial as failed
-                optimization_config = self.experiment.optimization_config
-                if (
-                    optimization_config is not None
-                    and metric_name in optimization_config.metric_names
-                    and not self.experiment.metrics[metric_name].is_recoverable_fetch_e(
-                        metric_fetch_e=metric_fetch_e
-                    )
-                ):
-                    status = self._mark_err_trial_status(
-                        trial=self.experiment.trials[trial_index],
-                        metric_name=metric_name,
-                        metric_fetch_e=metric_fetch_e,
-                    )
-                    self.logger.warning(
-                        f"MetricFetchE INFO: Because {metric_name} is an objective, "
-                        f"marking trial {trial_index} as {status}."
-                    )
-                    self._num_trials_bad_due_to_err += 1
-                    continue
-
                 self.logger.info(
-                    "MetricFetchE INFO: Continuing optimization even though "
-                    "MetricFetchE encountered."
+                    "MetricFetchE INFO: Trial status remains unchanged; metric data "
+                    "availability is tracked separately."
                 )
                 continue
 
@@ -2163,22 +2173,11 @@ def _fetch_and_process_trials_data_results(
     def _report_metric_fetch_e(
         self,
         trial: BaseTrial,
-        metric_name: str,
+        metric_signature: str,
         metric_fetch_e: MetricFetchE,
     ) -> None:
         pass
 
-    def _mark_err_trial_status(
-        self,
-        trial: BaseTrial,
-        metric_name: str | None = None,
-        metric_fetch_e: MetricFetchE | None = None,
-    ) -> TrialStatus:
-        trial.mark_abandoned(
-            reason=metric_fetch_e.message if metric_fetch_e else None, unsafe=True
-        )
-        return TrialStatus.ABANDONED
-
     def _get_failure_rate_exceeded_error(
         self,
         num_bad_in_orchestrator: int,
@@ -2187,7 +2186,7 @@ def _get_failure_rate_exceeded_error(
         return FailureRateExceededError(
             (
                 f"{METRIC_FETCH_ERR_MESSAGE}\n"
-                if self._num_trials_bad_due_to_err > num_bad_in_orchestrator / 2
+                if self._num_trials_bad_due_to_err > 0
                 else ""
             )
             + " Original error message: "