Update BenchmarkRunner to mark trials with NaN data as ABANDONED

Summary: Update `BenchmarkRunner.poll_trial_status` to mark trials with NaN data as `ABANDONED`.

Differential Revision: D87094541

Author: Sunny Shen

ax/benchmark/benchmark_runner.py

@@ -5,6 +5,7 @@
 
 # pyre-strict
 
+import logging
 import warnings
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass, field
@@ -26,6 +27,8 @@
 from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry
 from ax.utils.testing.backend_simulator import BackendSimulator, BackendSimulatorOptions
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 def _dict_of_arrays_to_df(
     Y_true_by_arm: Mapping[str, npt.NDArray],
@@ -264,12 +267,60 @@ def run(self, trial: BaseTrial) -> dict[str, BenchmarkTrialMetadata]:
         )
         return {"benchmark_metadata": metadata}
 
+    def _trial_has_nan(self, trial: BaseTrial) -> bool:
+        """Check if a trial's benchmark metadata contains NaN in the mean column.
+
+        Args:
+            trial: The trial to check.
+
+        Returns:
+            True if any metric's DataFrame has NaN in the "mean" column.
+        """
+        run_metadata = trial.run_metadata
+        if not isinstance(run_metadata, dict):
+            return False
+        metadata = run_metadata.get("benchmark_metadata")
+        if metadata is None:
+            return False
+        for df in metadata.dfs.values():
+            if not df.empty and df["mean"].isna().any():
+                return True
+        return False
+
     def poll_trial_status(
         self, trials: Iterable[BaseTrial]
     ) -> dict[TrialStatus, set[int]]:
         if self.simulated_backend_runner is None:
-            return {TrialStatus.COMPLETED: {t.index for t in trials}}
-        return self.simulated_backend_runner.poll_trial_status(trials=trials)
+            statuses: dict[TrialStatus, set[int]] = {
+                TrialStatus.COMPLETED: {t.index for t in trials}
+            }
+        else:
+            statuses = self.simulated_backend_runner.poll_trial_status(trials=trials)
+
+        # Move completed trials with NaN data to ABANDONED.
+        completed_indices = statuses.get(TrialStatus.COMPLETED, set())
+        if completed_indices:
+            trials_by_index = {t.index: t for t in trials}
+            nan_indices = {
+                idx
+                for idx in completed_indices
+                if (trial := trials_by_index.get(idx)) is not None
+                and self._trial_has_nan(trial)
+            }
+            if nan_indices:
+                for idx in nan_indices:
+                    logger.info(
+                        f"Trial {idx} has NaN in metrics and will be "
+                        "marked as ABANDONED."
+                    )
+                statuses.pop(TrialStatus.COMPLETED)
+                remaining = completed_indices - nan_indices
+                if remaining:
+                    statuses[TrialStatus.COMPLETED] = remaining
+                statuses[TrialStatus.ABANDONED] = (
+                    statuses.get(TrialStatus.ABANDONED, set()) | nan_indices
+                )
+        return statuses
 
     @classmethod
     def serialize_init_args(cls, obj: Any) -> dict[str, Any]:

ax/benchmark/tests/test_benchmark_runner.py

@@ -445,3 +445,68 @@ def test_heterogeneous_step_runtime(self) -> None:
                 ValueError, "Step duration must be non-negative"
             ):
                 runner.run(trial=trial)
+
+    def test_nan_trials_abandoned(self) -> None:
+        """Trials with NaN data are marked ABANDONED, valid trials COMPLETED."""
+        test_function = IdentityTestFunction(outcome_names=["objective"])
+
+        with self.subTest("without simulated backend"):
+            runner = BenchmarkRunner(
+                test_function=test_function,
+                noise=GaussianNoise(noise_std=0.0),
+            )
+            experiment = Experiment(
+                name="test_nan_sync",
+                is_test=True,
+                runner=runner,
+                search_space=Mock(spec=SearchSpace),
+            )
+            trial_valid = Trial(experiment=experiment)
+            trial_valid.add_arm(Arm(name="0_0", parameters={"x0": 1.5}))
+            trial_valid.run()
+
+            trial_nan = Trial(experiment=experiment)
+            trial_nan.add_arm(Arm(name="1_0", parameters={"x0": float("nan")}))
+            trial_nan.run()
+
+            trial_valid2 = Trial(experiment=experiment)
+            trial_valid2.add_arm(Arm(name="2_0", parameters={"x0": 2.0}))
+            trial_valid2.run()
+
+            statuses = runner.poll_trial_status([trial_valid, trial_nan, trial_valid2])
+            self.assertEqual(
+                statuses[TrialStatus.COMPLETED],
+                {trial_valid.index, trial_valid2.index},
+            )
+            self.assertEqual(statuses[TrialStatus.ABANDONED], {trial_nan.index})
+
+        with self.subTest("with simulated backend"):
+            runner = BenchmarkRunner(
+                test_function=test_function,
+                noise=GaussianNoise(noise_std=0.0),
+                force_use_simulated_backend=True,
+            )
+            experiment = Experiment(
+                name="test_nan_async",
+                is_test=True,
+                runner=runner,
+                search_space=Mock(spec=SearchSpace),
+            )
+            trial_nan = Trial(experiment=experiment)
+            trial_nan.add_arm(Arm(name="0_0", parameters={"x0": float("nan")}))
+            trial_nan.run()
+
+            trial_valid = Trial(experiment=experiment)
+            trial_valid.add_arm(Arm(name="1_0", parameters={"x0": 3.0}))
+            trial_valid.run()
+
+            # Advance simulated time so trials complete
+            simulator = none_throws(
+                none_throws(runner.simulated_backend_runner).simulator
+            )
+            while simulator.num_completed < 2:
+                simulator.update()
+
+            statuses = runner.poll_trial_status([trial_nan, trial_valid])
+            self.assertEqual(statuses[TrialStatus.COMPLETED], {trial_valid.index})
+            self.assertEqual(statuses[TrialStatus.ABANDONED], {trial_nan.index})