From 8a51f434e1950a7927a1492504641423fcdac389 Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Fri, 18 Sep 2026 08:15:28 -0700 Subject: [PATCH] Decouple metric fetch errors from trial status in Orchestrator (#5287) Summary: 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: esantorella Differential Revision: D119207750 --- ax/core/metric.py | 10 +- ax/orchestration/orchestrator.py | 175 +++++++++---------- ax/orchestration/tests/test_orchestrator.py | 184 +++++++------------- 3 files changed, 157 insertions(+), 212 deletions(-) diff --git a/ax/core/metric.py b/ax/core/metric.py index 24240dfe4b6..137d9d0ae42 100644 --- a/ax/core/metric.py +++ b/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( diff --git a/ax/orchestration/orchestrator.py b/ax/orchestration/orchestrator.py index bce2bab7bac..2c92a4a1cda 100644 --- a/ax/orchestration/orchestrator.py +++ b/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: " diff --git a/ax/orchestration/tests/test_orchestrator.py b/ax/orchestration/tests/test_orchestrator.py index fdbf0c8143c..8a14c6244af 100644 --- a/ax/orchestration/tests/test_orchestrator.py +++ b/ax/orchestration/tests/test_orchestrator.py @@ -28,7 +28,7 @@ from ax.core.experiment import Experiment from ax.core.experiment_status import ExperimentStatus from ax.core.generator_run import GeneratorRun -from ax.core.metric import Metric +from ax.core.metric import Metric, MetricFetchE from ax.core.multi_type_experiment import MultiTypeExperiment from ax.core.objective import Objective from ax.core.observation import ObservationFeatures @@ -40,12 +40,7 @@ get_pending_observation_features_based_on_trial_status, ) from ax.early_stopping.strategies import BaseEarlyStoppingStrategy -from ax.exceptions.core import ( - AxError, - OptimizationComplete, - UnsupportedError, - UserInputError, -) +from ax.exceptions.core import OptimizationComplete, UnsupportedError, UserInputError from ax.exceptions.generation_strategy import AxGenerationException from ax.generation_strategy.dispatch_utils import choose_generation_strategy_legacy from ax.generation_strategy.generation_strategy import ( @@ -99,6 +94,7 @@ from ax.storage.sqa_store.with_db_settings_base import WithDBSettingsBase from ax.utils.common.constants import Keys from ax.utils.common.logger import AX_ROOT_LOGGER_NAME +from ax.utils.common.result import Err from ax.utils.common.testutils import TestCase from ax.utils.common.timeutils import current_timestamp_in_millis from ax.utils.testing.core_stubs import ( @@ -1015,6 +1011,36 @@ def test_failure_rate_all_failed(self) -> None: orchestrator.run_all_trials() self.assertEqual(len(orchestrator.experiment.trials), 2) + def test_failure_rate_error_uses_trigger_counts_for_incomplete_data(self) -> None: + orchestrator = Orchestrator( + experiment=self.branin_experiment, + generation_strategy=self.sobol_GS_no_parallelism, + options=OrchestratorOptions( + tolerated_trial_failure_rate=0.2, + **self.orchestrator_options_kwargs, + ), + db_settings=self.db_settings_if_always_needed, + ) + trial = self.branin_experiment.new_trial() + trial.mark_running(no_runner_required=True).mark_completed() + + with self.assertRaisesRegex( + FailureRateExceededError, + "at least 1 out of the first 1 trials", + ): + orchestrator.error_if_failure_rate_exceeded(force_check=True) + + later_trial = self.branin_experiment.new_trial() + later_trial.mark_running(no_runner_required=True).mark_completed() + with self.assertRaisesRegex( + FailureRateExceededError, + "at least 1 out of the first 1 trials", + ): + orchestrator.error_if_failure_rate_exceeded(force_check=True) + + self.assertEqual(trial.status, TrialStatus.COMPLETED) + self.assertEqual(later_trial.status, TrialStatus.COMPLETED) + def test_sqa_storage_without_experiment_name(self) -> None: init_test_engine_and_session_factory(force_init=True) gs = self.two_sobol_steps_GS @@ -1912,10 +1938,11 @@ def test_poll_trial_status_abandons_trial_on_individual_failure(self) -> None: ) ) - def test_fetch_and_process_trials_data_results_failed_objective_available_while_running( # noqa - self, - ) -> None: + def test_fetch_error_for_running_metric_with_distinct_signature(self) -> None: gs = self.two_sobol_steps_GS + self.branin_timestamp_map_metric_experiment.metrics[ + "branin_map" + ].signature_override = "branin_map_signature" with ( patch( f"{BraninTimestampMapMetric.__module__}.BraninTimestampMapMetric.f", @@ -1947,6 +1974,12 @@ def test_fetch_and_process_trials_data_results_failed_objective_available_while_ self.assertTrue( any("Waiting for completed trials" in msg for msg in lg.output) ) + logs = "\n".join(lg.output) + self.assertIn( + "continuing after failing to fetch branin_map_signature and retrying", + logs, + ) + self.assertNotIn("Failed to fetch branin_map_signature for trial 0", logs) self.assertEqual( orchestrator.experiment.trials[0].status, TrialStatus.COMPLETED ) @@ -1983,106 +2016,37 @@ def test_fetch_and_process_trials_data_results_failed_non_objective( orchestrator.experiment.trials[0].status, TrialStatus.COMPLETED ) - def test_fetch_and_process_trials_data_results_failed_objective(self) -> None: - gs = self.two_sobol_steps_GS + def test_fetch_and_process_trials_data_results_unregistered_metric(self) -> None: orchestrator = Orchestrator( experiment=self.branin_experiment, - generation_strategy=gs, - options=OrchestratorOptions( - **self.orchestrator_options_kwargs, - ), + generation_strategy=self.two_sobol_steps_GS, + options=OrchestratorOptions(**self.orchestrator_options_kwargs), db_settings=self.db_settings_if_always_needed, ) - with ( - patch( - f"{BraninMetric.__module__}.BraninMetric.f", - side_effect=Exception("yikes!"), - ), - patch( - f"{BraninMetric.__module__}.BraninMetric.is_available_while_running", - return_value=False, - ), - self.assertLogs(logger="ax.orchestration.orchestrator") as lg, - ): - # This trial will fail - with self.assertRaises(FailureRateExceededError): - orchestrator.run_n_trials(max_trials=1) - self.assertTrue( - any( - re.search(r"Failed to fetch (branin|m1) for trial 0", warning) - is not None - for warning in lg.output - ) - ) - self.assertTrue( - any( - re.search( - r"Because (branin|m1) is an objective, marking trial 0 as " - "TrialStatus.ABANDONED", - warning, + trial = self.branin_experiment.new_trial() + trial.mark_running(no_runner_required=True).mark_completed() + results = { + trial.index: { + "collection_member": Err( + MetricFetchE(message="fetch failed", exception=None) ) - is not None - for warning in lg.output - ) - ) - self.assertEqual( - orchestrator.experiment.trials[0].status, TrialStatus.ABANDONED - ) + } + } - def test_fetch_and_process_trials_data_results_failed_objective_but_recoverable( - self, - ) -> None: - gs = self.two_sobol_steps_GS - orchestrator = Orchestrator( - experiment=self.branin_experiment, - generation_strategy=gs, - options=OrchestratorOptions( - enforce_immutable_search_space_and_opt_config=False, - **self.orchestrator_options_kwargs, - ), - db_settings=self.db_settings_if_always_needed, - ) - BraninMetric.recoverable_exceptions = {AxError, TypeError} - # we're throwing a recoverable exception because UserInputError - # is a subclass of AxError - with ( - patch( - f"{BraninMetric.__module__}.BraninMetric.f", - side_effect=UserInputError("yikes!"), - ), - patch( - f"{BraninMetric.__module__}.BraninMetric.is_available_while_running", - return_value=False, - ), - self.assertLogs(logger="ax.orchestration.orchestrator") as lg, + with patch.object( + self.branin_experiment, + "fetch_trials_data_results", + return_value=results, ): - orchestrator.run_n_trials(max_trials=1) - self.assertTrue( - any( - re.search(r"Failed to fetch (branin|m1) for trial 0", warning) - is not None - for warning in lg.output - ), - lg.output, - ) - self.assertTrue( - any( - re.search( - "MetricFetchE INFO: Continuing optimization even though " - "MetricFetchE encountered", - warning, - ) - is not None - for warning in lg.output + actual = orchestrator._fetch_and_process_trials_data_results( + trial_indices=[trial.index] ) - ) - self.assertEqual( - orchestrator.experiment.trials[0].status, TrialStatus.COMPLETED - ) - def test_fetch_and_process_trials_data_results_failed_objective_not_recoverable( - self, - ) -> None: + self.assertEqual(actual, results) + self.assertEqual(trial.status, TrialStatus.COMPLETED) + self.assertEqual(orchestrator._num_metric_fetch_e_encountered, 1) + + def test_fetch_error_does_not_change_completed_trial_status(self) -> None: gs = self.two_sobol_steps_GS orchestrator = Orchestrator( experiment=self.branin_experiment, @@ -2092,9 +2056,6 @@ def test_fetch_and_process_trials_data_results_failed_objective_not_recoverable( ), db_settings=self.db_settings_if_always_needed, ) - # we're throwing a unrecoverable exception because Exception is not subclass - # of either error type in recoverable_exceptions - BraninMetric.recoverable_exceptions = {AxError, TypeError} with ( patch( f"{BraninMetric.__module__}.BraninMetric.f", @@ -2106,7 +2067,6 @@ def test_fetch_and_process_trials_data_results_failed_objective_not_recoverable( ), self.assertLogs(logger="ax.orchestration.orchestrator") as lg, ): - # This trial will fail with self.assertRaises(FailureRateExceededError): orchestrator.run_n_trials(max_trials=1) self.assertTrue( @@ -2116,20 +2076,10 @@ def test_fetch_and_process_trials_data_results_failed_objective_not_recoverable( for warning in lg.output ) ) - self.assertTrue( - any( - re.search( - r"Because (branin|m1) is an objective, marking trial 0 as " - "TrialStatus.ABANDONED", - warning, - ) - is not None - for warning in lg.output - ) - ) self.assertEqual( - orchestrator.experiment.trials[0].status, TrialStatus.ABANDONED + orchestrator.experiment.trials[0].status, TrialStatus.COMPLETED ) + self.assertEqual(orchestrator._num_trials_bad_due_to_err, 1) def test_should_consider_optimization_complete(self) -> None: # Tests non-GSS parts of the completion criterion.