diff --git a/ax/benchmark/benchmark_metric.py b/ax/benchmark/benchmark_metric.py index e27fa586444..e4bfeb85169 100644 --- a/ax/benchmark/benchmark_metric.py +++ b/ax/benchmark/benchmark_metric.py @@ -3,8 +3,6 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -# pyre-strict - """ Metric classes for Ax benchmarking. @@ -183,7 +181,7 @@ def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult start_time = none_throws(sim_trial.sim_start_time) if sim_trial.sim_completed_time is None: # Still running - max_t = backend_simulator.time - start_time + max_t = backend_simulator.time elif sim_trial.sim_completed_time > backend_simulator.time: raise RuntimeError( "The trial's completion time is in the future! This is " @@ -192,10 +190,10 @@ def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> MetricFetchResult ) else: # Completed, may have stopped early -- can't assume all data available - completed_time = none_throws(sim_trial.sim_completed_time) - max_t = completed_time - start_time + max_t = none_throws(sim_trial.sim_completed_time) - available_data = df[df["virtual runtime"] <= max_t] + # Compare absolute times to avoid cancellation when subtracting the start. + available_data = df[start_time + df["virtual runtime"] <= max_t] if not self.observe_noise_sd: available_data.loc[:, "sem"] = np.nan diff --git a/ax/benchmark/tests/test_benchmark_metric.py b/ax/benchmark/tests/test_benchmark_metric.py index c6b46321fa3..0ccddf71cb0 100644 --- a/ax/benchmark/tests/test_benchmark_metric.py +++ b/ax/benchmark/tests/test_benchmark_metric.py @@ -3,8 +3,6 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -# pyre-strict - import math from itertools import product from unittest.mock import Mock @@ -347,6 +345,59 @@ def test_fetch_trial_multiple_time_steps_with_simulator(self) -> None: self._test_fetch_trial_multiple_time_steps_with_simulator(batch=False) self._test_fetch_trial_multiple_time_steps_with_simulator(batch=True) + def test_fetch_trial_data_at_simulated_time_boundary(self) -> None: + start_time = 192.0 + runtime = 320.10965660870767 + completed_time = start_time + runtime + self.assertLess(completed_time - start_time, runtime) + for early_stop, metrics in product( + [False, True], [self.map_metrics, self.tv_metrics] + ): + metric = metrics["test_metric1"] + with self.subTest(early_stop=early_stop, metric_cls=type(metric)): + trial = get_test_trial(multiple_time_steps=True) + metadata = trial.run_metadata["benchmark_metadata"] + df = metadata.dfs["test_metric1"] + df["virtual runtime"] = [0.0, runtime / 2, runtime] + df["mean"] = [0.0, 1.0, 2.0] + simulator = BackendSimulator( + options=BackendSimulatorOptions(internal_clock=start_time), + verbose_logging=False, + ) + simulator.run_trial(trial.index, runtime=runtime) + trial.update_run_metadata( + { + "benchmark_metadata": BenchmarkTrialMetadata( + dfs=metadata.dfs, backend_simulator=simulator + ) + } + ) + simulator.options.internal_clock = np.nextafter(completed_time, -np.inf) + data = assert_is_instance(metric.fetch_trial_data(trial).value, Data) + self.assertEqual(data.df["mean"].tolist(), [1.0]) + if isinstance(metric, MapMetric): + self.assertEqual(data.full_df["step"].tolist(), [0, 1]) + if early_stop: + simulator.stop_trial(trial.index) + else: + simulator.options.internal_clock = completed_time + simulator._update_completed(completed_time) + data = assert_is_instance(metric.fetch_trial_data(trial).value, Data) + self.assertEqual(data.df["mean"].tolist(), [1.0 if early_stop else 2.0]) + if isinstance(metric, MapMetric): + self.assertEqual( + data.full_df["step"].tolist(), + [0, 1] if early_stop else [0, 1, 2], + ) + simulator.options.internal_clock = completed_time + 1 + data = assert_is_instance(metric.fetch_trial_data(trial).value, Data) + self.assertEqual(data.df["mean"].tolist(), [1.0 if early_stop else 2.0]) + if isinstance(metric, MapMetric): + self.assertEqual( + data.full_df["step"].tolist(), + [0, 1] if early_stop else [0, 1, 2], + ) + def test_sim_trial_completes_in_future_raises(self) -> None: simulator = BackendSimulator() simulator.run_trial(trial_index=0, runtime=0)