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.