Skip to content

Inference trace and Best Point Recommendation (BPR) bugfix #4128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+312 −257
Open

Inference trace and Best Point Recommendation (BPR) bugfix #4128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 11 additions & 8 deletions ax/benchmark/benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,10 @@
from ax.core.utils import get_model_times
from ax.generation_strategy.generation_strategy import GenerationStrategy
from ax.service.orchestrator import Orchestrator
from ax.service.utils.best_point import get_trace
from ax.service.utils.best_point_mixin import BestPointMixin
from ax.service.utils.best_point import (
get_best_parameters_from_model_predictions_with_trial_index,
get_trace,
)
from ax.service.utils.orchestrator_options import OrchestratorOptions, TrialType
from ax.utils.common.logger import DEFAULT_LOG_LEVEL, get_logger
from ax.utils.common.random import with_rng_seed
Expand Down Expand Up @@ -326,9 +328,8 @@ def get_best_parameters(
trial_indices: Iterable[int] | None = None,
) -> TParameterization | None:
"""
Get the most promising point.

Only SOO is supported. It will return None if no best point can be found.
Get the most promising point. Returns None if no point is predicted to
satisfy all outcome constraints.

Args:
experiment: The experiment to get the data from. This should contain
Expand All @@ -338,9 +339,9 @@ def get_best_parameters(
best point.
trial_indices: Use data from only these trials. If None, use all data.
"""
result = BestPointMixin._get_best_trial(
result = get_best_parameters_from_model_predictions_with_trial_index(
experiment=experiment,
generation_strategy=generation_strategy,
adapter=generation_strategy.adapter,
trial_indices=trial_indices,
)
if result is None:
Expand Down Expand Up @@ -507,7 +508,7 @@ def run_optimization_with_orchestrator(

orchestrator = Orchestrator(
experiment=experiment,
generation_strategy=method.generation_strategy.clone_reset(),
generation_strategy=method.generation_strategy,
options=orchestrator_options,
)

Expand Down Expand Up @@ -562,6 +563,8 @@ def benchmark_replication(
Return:
``BenchmarkResult`` object.
"""
# Reset the generation strategy to ensure that it is in an unused state.
method.generation_strategy = method.generation_strategy.clone_reset()
experiment = run_optimization_with_orchestrator(
problem=problem,
method=method,
Expand Down
97 changes: 23 additions & 74 deletions ax/benchmark/tests/test_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

import numpy as np
import torch
from ax.adapter.factory import get_sobol
from ax.adapter.registry import Generators
from ax.benchmark.benchmark import (
_get_oracle_value_of_params,
Expand All @@ -38,7 +37,6 @@
from ax.benchmark.benchmark_problem import (
BenchmarkProblem,
create_problem_from_botorch,
get_continuous_search_space,
get_moo_opt_config,
get_soo_opt_config,
)
Expand Down Expand Up @@ -275,8 +273,6 @@ def _test_replication_async(self, map_data: bool) -> None:
map_data: If True, the test function produces time-series data with
just one step, so behavior is the same as when map_data=False.
"""
method = get_async_benchmark_method()

complete_out_of_order_runtimes = {
0: 2,
1: 1,
Expand Down Expand Up @@ -345,6 +341,8 @@ def _test_replication_async(self, map_data: bool) -> None:
}

for case_name, step_runtime_fn in step_runtime_fns.items():
method = get_async_benchmark_method()

with self.subTest(case_name, step_runtime_fn=step_runtime_fn):
problem = get_async_benchmark_problem(
map_data=map_data,
Expand Down Expand Up @@ -417,12 +415,9 @@ def _test_replication_async(self, map_data: bool) -> None:
},
f"Failure for trial {trial_index} with {case_name}",
)
self.assertFalse(np.isnan(result.inference_trace).any())
self.assertEqual(
result.inference_trace.tolist(),
expected_traces[case_name],
msg=case_name,
)
# Evaluating inference trace here is not relevant since the GS
# is not model-based.
self.assertTrue(np.isnan(result.inference_trace).all())
self.assertEqual(
result.oracle_trace.tolist(),
expected_traces[case_name],
Expand Down Expand Up @@ -491,6 +486,7 @@ def test_run_optimization_with_orchestrator(self) -> None:
none_throws(runner.simulated_backend_runner).simulator._verbose_logging
)

method = get_async_benchmark_method()
with self.subTest("Logs not produced by default"), self.assertNoLogs(
level=logging.INFO, logger=logger
), self.assertNoLogs(logger=logger):
Expand Down Expand Up @@ -618,9 +614,9 @@ def test_early_stopping(self) -> None:
self.assertEqual(max_run, {0: 4, 1: 2, 2: 2, 3: 2})

def test_replication_variable_runtime(self) -> None:
method = get_async_benchmark_method(max_pending_trials=1)
for map_data in [False, True]:
with self.subTest(map_data=map_data):
method = get_async_benchmark_method(max_pending_trials=1)
problem = get_async_benchmark_problem(
map_data=map_data,
step_runtime_fn=lambda params: params["x0"] + 1,
Expand Down Expand Up @@ -793,7 +789,11 @@ def test_replication_mbm(self) -> None:
acquisition_cls=qLogNoisyExpectedImprovement,
distribute_replications=False,
),
get_augmented_branin_problem(fidelity_or_task="fidelity"),
get_single_objective_benchmark_problem(
observe_noise_sd=False,
num_trials=6,
report_inference_value_as_trace=True,
),
"MBM::SingleTaskGP_qLogNEI",
),
]:
Expand Down Expand Up @@ -1196,6 +1196,7 @@ def test_get_opt_trace_by_cumulative_epochs(self) -> None:
):
get_opt_trace_by_steps(experiment=experiment)

method = get_sobol_benchmark_method(distribute_replications=False)
with self.subTest("Constrained"):
problem = get_benchmark_problem("constrained_gramacy_observed_noise")
experiment = self.run_optimization_with_orchestrator(
Expand Down Expand Up @@ -1237,72 +1238,20 @@ def test_get_benchmark_result_with_cumulative_steps(self) -> None:
self.assertLessEqual(transformed.score_trace.min(), result.score_trace.min())

def test_get_best_parameters(self) -> None:
"""
Whether this produces the correct values is tested more thoroughly in
other tests such as `test_replication_with_inference_value` and
`test_get_inference_trace_from_params`. Setting up an experiment with
data and trials without just running a benchmark is a pain, so in those
tests, we just run a benchmark.
"""
gs = get_sobol_generation_strategy()

search_space = get_continuous_search_space(bounds=[(0, 1)])
moo_config = get_moo_opt_config(outcome_names=["a", "b"], ref_point=[0, 0])
experiment = Experiment(
name="test",
is_test=True,
search_space=search_space,
optimization_config=moo_config,
experiment = get_experiment_with_observations(observations=[[1.0, 2.0]])
generation_strategy = get_sobol_generation_strategy()
mock_function = (
"ax.benchmark.benchmark."
"get_best_parameters_from_model_predictions_with_trial_index"
)

with self.subTest("MOO not supported"), self.assertRaisesRegex(
NotImplementedError, "Please use `get_pareto_optimal_parameters`"
):
get_best_parameters(experiment=experiment, generation_strategy=gs)

soo_config = get_soo_opt_config(outcome_names=["a"])
with self.subTest("Empty experiment"):
result = get_best_parameters(
experiment=experiment.clone_with(optimization_config=soo_config),
generation_strategy=gs,
)
with patch(mock_function, return_value=None):
result = get_best_parameters(experiment, generation_strategy)
self.assertIsNone(result)

with self.subTest("All constraints violated"):
experiment = get_experiment_with_observations(
observations=[[1, -1], [2, -1]],
constrained=True,
)
best_point = get_best_parameters(
experiment=experiment, generation_strategy=gs
)
self.assertIsNone(best_point)

with self.subTest("No completed trials"):
experiment = get_experiment_with_observations(observations=[])
sobol_generator = get_sobol(search_space=experiment.search_space)
for _ in range(3):
trial = experiment.new_trial(generator_run=sobol_generator.gen(n=1))
trial.run()
best_point = get_best_parameters(
experiment=experiment, generation_strategy=gs
)
self.assertIsNone(best_point)

experiment = get_experiment_with_observations(
observations=[[1], [2]], constrained=False
)
with self.subTest("Working case"):
best_point = get_best_parameters(
experiment=experiment, generation_strategy=gs
)
self.assertEqual(best_point, experiment.trials[1].arms[0].parameters)

with self.subTest("Trial indices"):
best_point = get_best_parameters(
experiment=experiment, generation_strategy=gs, trial_indices=[0]
)
self.assertEqual(best_point, experiment.trials[0].arms[0].parameters)
with patch(mock_function, return_value=(0, {"x": 1.0}, None)):
result = get_best_parameters(experiment, generation_strategy)
self.assertEqual(result, {"x": 1.0})

def test_get_benchmark_result_from_experiment_and_gs(self) -> None:
problem = get_single_objective_benchmark_problem()
Expand Down
Loading